pytorch Issue: SequentialImpl and AnyModule do not support custom user-defined Modules inherited from Module (JVM crash / compilation failure) and ModuleDict need mapping Insert method

[mullerhai]

Library: javacpp-presets/pytorch
Version: (Please fill in your used version)
Related C++ Library: LibTorch (fully supports this feature)
Related Python Library: PyTorch (fully supports this feature)
Problem Description
There are two critical functional defects in SequentialImpl and AnyModule when using custom user-defined layers that inherit from Module in Java, which are fully supported in both C++ LibTorch and Python PyTorch, but broken in JavaCPP PyTorch:

1. SequentialImpl lacks native support for inserting custom Module subclasses
   The C++ LibTorch SequentialImpl provides a variadic template constructor and push_back method that directly accepts any module derived from torch::nn::Module (including custom layers):

运行
// C++ LibTorch (fully supported)
template <typename... Modules>
explicit SequentialImpl(Modules&&... modules) {
    modules_.reserve(sizeof...(Modules));
    push_back(std::forward<Modules>(modules)...);
}
void push_back(Module module);

#include <torch/torch.h>
#include <iostream>

// 1. 自定义层：必须继承 torch::nn::Module
struct MyCustomLayerImpl : torch::nn::Module {
    // 核心：实现 forward 函数
    torch::Tensor forward(torch::Tensor x) {
        x = torch::relu(x);            // ReLU 激活
        x = x / (torch::max(x) + 1e-8f); // 归一化
        return x;
    }
};

// LibTorch 固定用法：创建模块智能指针
TORCH_MODULE(MyCustomLayer);

int main() {
    // 2. 直接把自定义层插入 Sequential！
    auto model = torch::nn::Sequential(
        torch::nn::Linear(10, 20),
        MyCustomLayer(),  // 自定义层
        torch::nn::Linear(20, 5)
    );

    // 测试
    torch::Tensor x = torch::randn({2, 10});
    torch::Tensor out = model->forward(x);
    std::cout << "输出形状: " << out.sizes() << std::endl;
    return 0;
}

# 自定义 Layer：必须继承 nn.Module
class MyCustomLayer(nn.Module):
    def __init__(self):
        super().__init__()  # 必须调用父类初始化

    # 必须实现 forward
    def forward(self, x):
        x = torch.relu(x)
        x = x / (torch.max(x) + 1e-8)  # 归一化
        return x

# 直接插入 Sequential！
model = nn.Sequential(
    nn.Linear(10, 20),
    MyCustomLayer(),  # 自定义层
    nn.Linear(20, 5)
)

# 测试
x = torch.randn(2, 10)
out = model(x)
print(out.shape)  # torch.Size([2, 5])

The Python PyTorch nn.Sequential also directly accepts custom nn.Module subclasses.
JavaCPP PyTorch:
SequentialImpl does not have a push_back(Module module) native method that directly accepts base Module type.
Cannot compile when directly adding a custom Module (e.g., seq.push_back(new Dice()) → compilation error).

2. AnyModule wrapper causes JVM crash at runtime for custom ModulesAs a workaround, users wrap custom modules with AnyModule (compiles successfully), but this causes an immediate JVM crash at runtime:java运行// Compiles, but JVM CRASH at runtimeseq.push_back(new AnyModule(new Dice()));This works perfectly for built-in modules (e.g., LinearImpl, ReLUImpl), but fails catastrophically for user-defined Module subclasses.
   Reproducible Code Example
3. Custom Module (inherits from Module)

运行
import org.bytedeco.javacpp.DoublePointer;
import org.bytedeco.pytorch.Module;
import org.bytedeco.pytorch.Tensor;
import static org.bytedeco.pytorch.global.torch.*;

// Custom layer inheriting from base Module (identical pattern to C++/Python)
public class Dice extends Module {
    private double epsilon;
    private Tensor alpha;

    public Dice(double epsilon) {
        super("Dice");
        this.epsilon = epsilon;
        this.alpha = register_parameter("alpha", randn(new long[]{1}));
    }

    public Dice() {
        this(1e-3);
    }

    public Tensor forward(Tensor x) {
        Tensor avg = x.mean(1).unsqueeze(1);
        Tensor diff = x.sub(avg);
        Tensor sq = diff.mul(diff);
        Tensor var = sq.sum(1).unsqueeze(1);
        var = var.add(tensor(new DoublePointer(epsilon)));
        Tensor ps = x.sub(avg).div(sqrt(var));
        ps = sigmoid(ps);
        Tensor part1 = ps.mul(x);
        Tensor part2 = ones_like(ps).mul(alpha).mul(ones_like(ps).sub(ps)).mul(x);
        return part1.add(part2);
    }
}

2. Test Code with Sequential

运行
import org.bytedeco.pytorch.AnyModule;
import org.bytedeco.pytorch.LinearImpl;
import org.bytedeco.pytorch.ReLUImpl;
import org.bytedeco.pytorch.SequentialImpl;

public class SequentialTest {
    public static void main(String[] args) {
        SequentialImpl seq = new SequentialImpl();

        // Built-in modules work fine
        seq.push_back(new AnyModule(new LinearImpl(128, 64)));
        ReLUImpl actModule = ActivationFactory.relu();
        seq.push_back(new AnyModule(actModule));

        // 1. DIRECT INSERT: Compilation FAILED
        // seq.push_back(new Dice());

        // 2. AnyModule WRAPPER: Compiles OK, but JVM CRASHES at runtime
        seq.push_back(new AnyModule(new Dice()));
    }
}

Expected Behavior
Consistent with C++ LibTorch and Python PyTorch:
SequentialImpl should expose a native method:

运行
public native void push_back(@ByVal Module module);

Allowing direct insertion of any custom subclass of Module without wrappers.
AnyModule should safely wrap and hold custom Module subclasses without JVM crashes.
Custom Module implementations should work identically to built-in modules inside Sequential.
Current Status
Direct insertion → Compilation error
AnyModule wrapper → Runtime JVM crash
C++/Python counterparts → Fully functional
This is a core functionality gap that blocks using custom neural network layers in JavaCPP PyTorch.

[mullerhai]

I think SequentialImpl and AnyModule do not support custom user-defined Modules inherited from Module maybe a bug block complex model constructed

[mullerhai]

Here need mapping before javacpp compile

I think SequentialImpl and AnyModule do not support custom user-defined Modules inherited from Module maybe a bug block complex model constructed

we need implement this in SequentialImpl

public native void push_back(@SharedPtr("torch::nn::Module") @ByVal Module module);

[mullerhai]

I also found one little bug ，the ModuleDictImpl contain update method ,but not contain 【insert】 method ，maybe just forget mapping ，the cpp source code found

  /// Private `OrderedDict` holding the key-Module pairs.
  torch::OrderedDict<std::string, std::shared_ptr<Module>> modules_;

  /// Insert a key-module pair by overwriting existing keys,
  /// and register or replace the `Module`.
  void insert(const std::string& key, std::shared_ptr<Module> module) {
    if (contains(key)) {
      modules_[key] = std::move(module);
      replace_module(key, modules_[key]);
    } else {
      modules_.insert(key, std::move(module));
      register_module(key, modules_.back().value());
    }
  }
};

[mullerhai]

Here need mapping before javacpp compile

I think SequentialImpl and AnyModule do not support custom user-defined Modules inherited from Module maybe a bug block complex model constructed

we need implement this in SequentialImpl

public native void push_back(@sharedptr("torch::nn::Module") @byval Module module);

@saudet could you mapping implement > public native void push_back(@sharedptr("torch::nn::Module") @byval Module module); for SequentialImpl in next 1.5.14 snapshot version?

[mullerhai]

HI， @saudet ，

have to emphasize this again today. I am developing over 50 recommendation system models, and more than 30 of them rely heavily on MLPs and various layer lists. If Sequential cannot accept user-defined Modules, it will become a major roadblock for our work. This issue needs to be taken seriously, and the corresponding type mapping must be implemented. Please note that AnyModule does not address this problem — it is intended for a completely different scenario.

Feature Request: Support custom user-defined Modules in Sequential via presets mapping
Labels: enhancement, pytorch, presets
Issue Content
Description
Currently, the Sequential implementation in JavaCPP PyTorch presets cannot accept arbitrary custom user-defined Module instances as child layers.
From the underlying implementation perspective, this capability only requires adding a simple javatext type mapping declaration in the infomap of PyTorch presets. Without this support, building even basic neural networks such as a standard MLP becomes extremely cumbersome on JVM languages (Java/Scala).
At present, developers have to manually use ArrayList to collect layers, iterate over the list, perform type checking for each element, and invoke the forward method one by one. This approach is verbose, error-prone, and breaks the native usage paradigm of PyTorch.
Python Example (Native & Concise Usage)
This is the standard and intuitive way to build an MLP with torch.nn.Sequential in Python:

运行
import torch
import torch.nn as nn

# Define custom layer
class CustomLinear(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.linear = nn.Linear(in_features, out_features)

    def forward(self, x):
        return torch.relu(self.linear(x))

# Use custom module directly in Sequential
model = nn.Sequential(
    CustomLinear(10, 32),
    nn.ReLU(),
    nn.Linear(32, 1)
)

# Forward pass
x = torch.randn(8, 10)
out = model(x)

class MLP(nn.Module):
    """Multi-layer perceptron with BN/activation/dropout per linear layer.

    Parameters
    ----------
    input_dim : int
        Input dimension of the first linear layer.
    output_layer : bool, default True
        If True, append a final Linear(*,1).
    dims : list, default []
        Hidden layer sizes.
    dropout : float, default 0
        Dropout probability.
    activation : str, default 'relu'
        Activation function (sigmoid, relu, prelu, dice, softmax).

    Shape
    -----
    Input: ``(B, input_dim)``  
    Output: ``(B, 1)`` or ``(B, dims[-1])``
    """

    def __init__(self, input_dim, output_layer=True, dims=None, dropout=0, activation="relu"):
        super().__init__()
        if dims is None:
            dims = []
        layers = list()
        for i_dim in dims:
            layers.append(nn.Linear(input_dim, i_dim))
            layers.append(nn.BatchNorm1d(i_dim))
            layers.append(activation_layer(activation))
            layers.append(nn.Dropout(p=dropout))
            input_dim = i_dim
        if output_layer:
            layers.append(nn.Linear(input_dim, 1))
        self.mlp = nn.Sequential(*layers)

    def forward(self, x):
        return self.mlp(x)

Scala Example (Current Workaround — Verbose & Clunky)
Due to the lack of support for custom Module inside Sequential, we have to implement manual collection, type check and iteration:

import org.bytedeco.pytorch._
import scala.collection.mutable.ArrayBuffer

/**
 * Multi-Layer Perceptron
 */
class MLP(
  inputDim: Long,
  hiddenDims: List[Long],
  outputDim: Long = 1,
  activation: String = "relu",
  dropout: Float = 0.0f,
  useBatchNorm: Boolean = false,
  useLayerNorm: Boolean = false,
  device: String = DeviceSupport.backend
) extends Module {

  private val layers = mutable.ListBuffer[Module]()
  private var prevDim = inputDim

  // Build MLP layers
  hiddenDims.foreach { dim =>
    val linear = new LinearImpl(prevDim, dim)
    register_module(s"linear_${layers.size}", linear)
    layers += linear

    if (useLayerNorm) {
      val ln = new LayerNormImpl(new LongVector(dim))
      register_module(s"ln_${layers.size}", ln)
      layers += ln
    } else if (useBatchNorm && !activation.equals("relu")) {
      val bn = new BatchNorm1dImpl(new BatchNormOptions(dim))
      register_module(s"bn_${layers.size}", bn)
      layers += bn
    }

    layers += createActivation(activation)

    if (dropout > 0) {
      val drop = new DropoutImpl(dropout)
      register_module(s"dropout_${layers.size}", drop)
      layers += drop
    }

    prevDim = dim
  }

  // Output layer
  val outputLinear = new LinearImpl(prevDim, outputDim)
  register_module("output", outputLinear)
  layers += outputLinear

  // Move all layers to target device
  if (device != "cpu") {
    val dev = new org.bytedeco.pytorch.Device(device)
    layers.foreach { m =>
      m.to(dev, false)
    }
    this.to(dev, false)
  }

  def forward(x: Tensor): Tensor = {
    var result: Tensor = x
    layers.foreach { layer =>
      result = layer match {
        case mlp: LinearImpl => mlp.forward(result)
        case bn: BatchNorm1dImpl => bn.forward(result)
        case ln: LayerNormImpl => ln.forward(result)
        case act: ReLUImpl => act.forward(result)
        case act: SigmoidImpl => act.forward(result)
        case act: TanhImpl => act.forward(result)
        case act: SiLUImpl => act.forward(result)
        case act: GELUImpl => act.forward(result)
        case drop: DropoutImpl => drop.forward(result)
        case _ => result
      }
    }
    result
  }

  private def createActivation(act: String): Module = {
    act.toLowerCase match {
      case "relu" => new ReLUImpl()
      case "sigmoid" => new SigmoidImpl()
      case "tanh" => new TanhImpl()
      case "silu" | "swish" => new SiLUImpl()
      case "gelu" => new GELUImpl()
      case "prelu" => new PReLUImpl()
      case "leaky_relu" | "leakyrelu" => new LeakyReLUImpl()
      case "none" | "identity" => new IdentityImpl()
      case _ => new ReLUImpl()
    }
  }
}



// Custom user-defined Module
class CustomLinear(inDim: Long, outDim: Long) extends Module {
  private val linear = new Linear(inDim, outDim)

  override def forward(x: Tensor): Tensor = {
    torch.relu(linear.forward(x))
  }
}

// Cannot add CustomLinear to Sequential directly, use ArrayList as workaround
val layerList = new ArrayBuffer[Module]()
layerList += new CustomLinear(10, 32)
layerList += new ReLU()
layerList += new Linear(32, 1)

// Manual iterate + type check + forward call
def modelForward(x: Tensor): Tensor = {
  var res = x
  for (layer <- layerList) {
    res = layer match {
      case l: CustomLinear => l.forward(res)
      case l: ReLU => l.forward(res)
      case l: Linear => l.forward(res)
      case _ => throw new IllegalArgumentException("Unknown layer type")
    }
  }
  res
}

// Run forward
val x = torch.randn(8, 10)
val output = modelForward(x)

Expected Behavior
After adding the corresponding javatext type mapping in presets infomap:
Allow directly adding any custom subclass of Module into Sequential constructor / add() method.
Keep the same coding style as native PyTorch, eliminate manual list management and type branching.
Request
Please add the necessary type mapping in PyTorch presets infomap to enable Sequential to accept arbitrary custom Module instances. This is a lightweight change but greatly improves usability for JVM developers using JavaCPP PyTorch.
Thank you!

[mullerhai]

HI, @saudet

now SequentialImpl push_back. user defined Module could work

  /**
   * Adds an existing {@link Module} (any subclass — built-in or
   * user-defined) to this {@code Sequential} container. The
   * actual C++ type is recovered via the typed {@code as<Impl>()}
   * helpers and the matching typed {@code push_back<Impl>}
   * overload is invoked. Throws when the underlying C++
   * type is not one of the recognised built-in layers.
   */
  public void push_back(@SharedPtr @ByVal Module module) {
      ModuleAsHelper.pushBackModule(this, module, null);
  }
  public void push_back(String name, @SharedPtr @ByVal Module module) {
      ModuleAsHelper.pushBackModule(this, module, name);
  }

package org.bytedeco.pytorch;

import java.lang.reflect.Constructor;
import java.lang.reflect.Method;
import org.bytedeco.javacpp.Pointer;

/**
 * Small reflection helper backing the generic
 * {@link Module#as(Class) Module.as(Class)} method and the generic
 * {@link SequentialImpl#push_back(Module)} entry point.
 *
 * <p>Both rely on the same idea: the C++ side does not let us instantiate
 * the {@code push_back<Module>} template because the base {@code Module}
 * class has no {@code forward()} method (see the
 * {@code has_forward<Module>} static_assert in {@code AnyModule}). We work
 * around that by performing the dispatch on the <em>Java</em> side via
 * the typed {@code asXxx()} methods JavaCPP already generates for the
 * 80+ built-in layers. The C++ JNI proxy is therefore only asked to
 * invoke instantiations of the template with concrete {@code *Impl} types
 * — the path the existing typed push_back overloads already exercise.
 *
 * <p>This class is internal and intentionally minimal — the goal is to
 * keep the byte-code injected by the preset as small as possible.
 */
final class ModuleAsHelper {

    private ModuleAsHelper() {}

    /**
     * Creates a fresh instance of {@code targetClass} wrapping
     * {@code module}'s native pointer. {@code targetClass} must declare a
     * public {@code (Pointer)} constructor.
     */
    static <T extends Module> T wrap(Module module, Class<T> targetClass) {
        try {
            Constructor<T> ctor = targetClass.getConstructor(Pointer.class);
            return ctor.newInstance(module);
        } catch (NoSuchMethodException e) {
            throw new IllegalArgumentException(
                    targetClass.getName()
                    + " must declare a public constructor accepting org.bytedeco.javacpp.Pointer",
                    e);
        } catch (ReflectiveOperationException e) {
            throw new RuntimeException(
                    "Failed to wrap pointer with " + targetClass.getName(), e);
        }
    }

    /**
     * Pushes {@code module} into {@code seq} by recovering its actual
     * concrete C++ type via the typed {@code as<Impl>()} helpers and
     * forwarding to the matching typed {@code push_back<Impl>()} overload.
     * The {@code name} parameter is optional — pass {@code null} to
     * auto-generate one (matching the C++ template's default behaviour).
     */
    static void pushBackModule(SequentialImpl seq, Module module, String name) {
        // Fast path: derive the asXxx method name from the simple class
        // name. No registry lookup, no array iteration — just one
        // reflection-based asXxx invocation.
        Class<?> moduleClass = module.getClass();
        String simpleName = moduleClass.getSimpleName();
        String asName = "as" + (simpleName.endsWith("Impl")
                ? simpleName.substring(0, simpleName.length() - "Impl".length())
                : simpleName);
        try {
            Method asMethod = Module.class.getMethod(asName);
            Object cast = asMethod.invoke(module);
            if (cast == null) {
                throw new IllegalArgumentException(
                        "torch::nn::push_back_module: module is not actually a "
                        + simpleName);
            }
            if (name == null) {
                Method pushBack = SequentialImpl.class.getMethod(
                        "push_back", moduleClass);
                pushBack.invoke(seq, cast);
            } else {
                Method pushBack = SequentialImpl.class.getMethod(
                        "push_back", String.class, moduleClass);
                pushBack.invoke(seq, name, cast);
            }
            return;
        } catch (NoSuchMethodException e) {
            // Fall through to the registered list below — this catches
            // custom user-defined classes that don't follow the
            // *Impl naming convention.
        } catch (ReflectiveOperationException e) {
            throw new RuntimeException(
                    "Failed to push_back " + module.getClass().getName(), e);
        }

        // Fallback path: walk the registered keys (only for unknown
        // user-defined classes).
        for (String key : keys) {
            try {
                Class<?> implClass = Class.forName("org.bytedeco.pytorch." + key);
                if (implClass.isInstance(module)) {
                    simpleName = key;
                    asName = "as" + (simpleName.endsWith("Impl")
                            ? simpleName.substring(0, simpleName.length() - "Impl".length())
                            : simpleName);
                    Method asMethod = Module.class.getMethod(asName);
                    Object cast = asMethod.invoke(module);
                    if (cast != null) {
                        if (name == null) {
                            Method pushBack = SequentialImpl.class.getMethod(
                                    "push_back", implClass);
                            pushBack.invoke(seq, cast);
                        } else {
                            Method pushBack = SequentialImpl.class.getMethod(
                                    "push_back", String.class, implClass);
                            pushBack.invoke(seq, name, cast);
                        }
                        return;
                    }
                }
            } catch (ClassNotFoundException | NoSuchMethodException ignored) {
                // skip
            } catch (ReflectiveOperationException e) {
                throw new RuntimeException(
                        "Failed to push_back " + module.getClass().getName(), e);
            }
        }
        throw new IllegalArgumentException(
                "torch::nn::push_back_module: unsupported module type: "
                + module.getClass().getName()
                + " — register the type with Module.registerCppType() first");
    }

    // Fallback list of known *Impl class simple names. Used only when
    // the fast path (module.getClass().getSimpleName()) doesn't match
    // a built-in layer name.
    private static final String[] keys = {
            "LinearImpl", "Conv1dImpl", "Conv2dImpl", "Conv3dImpl",
            "ConvTranspose1dImpl", "ConvTranspose2dImpl", "ConvTranspose3dImpl",
            "BatchNorm1dImpl", "BatchNorm2dImpl", "BatchNorm3dImpl",
            "ReLUImpl", "ReLU6Impl", "LeakyReLUImpl", "ELUImpl",
            "DropoutImpl", "Dropout2dImpl", "Dropout3dImpl",
            "FlattenImpl", "IdentityImpl", "EmbeddingImpl",
            "SequentialImpl", "ModuleListImpl", "ModuleDictImpl",
    };
}

[saudet]

Sure, please open a pull request

[mullerhai]

Sure, please open a pull request

I have make a pr，please review and check ，some detail maybe need you update the code ，thanks，first make pr to javacpp @saudet #1787

[mullerhai]

DropoutImpl must have bug ，insert to sequential will crash
private def buildResNetLike = {
// Kept linear-only for portability — same cast paths apply to conv
// modules; what matters for the benchmark is the generic dispatch.
val seq = new SequentialImpl
seq.push_back(new FlattenImpl)
seq.push_back(new LinearImpl(3L * 8L * 8L, 256L))
seq.push_back(new ReLUImpl)
// seq.push_back(new DropoutImpl(0.1));
seq.push_back(new LinearImpl(256L, 128L))
seq.push_back(new ReLUImpl)
seq.push_back(new LinearImpl(128L, 32L))
seq
}