ModuleListImpl erases concrete subclass types of stored Modules, causing ClassCastException when casting back to custom layer subclasses like MLP

[mullerhai]

Hi @saudet

Problem Description

When iterating elements from ModuleListImpl via its iterator / .get() method, all returned objects are exposed as the base type org.bytedeco.pytorch.Module with full type erasure at the JNI boundary.
Even though the underlying native module is our custom Scala layer (e.g. torchrec.basic.layers.MLP which extends Module), attempting to cast the returned Module reference back to the concrete subclass throws a hard ClassCastException.
Exception Stack
plaintext
Exception in thread "main" java.lang.ClassCastException: class org.bytedeco.pytorch.Module cannot be cast to class torchrec.basic.layers.MLP
(org.bytedeco.pytorch.Module and torchrec.basic.layers.MLP are in unnamed module of loader 'app')
Root Cause Analysis
ModuleListImpl JNI accessor only returns the generic base Module handle without retaining runtime subclass metadata of wrapped custom layers;
All custom layers (MLP, CapsuleNetwork, CrossLayer etc.) inherit from Module, but JNI layer type information is lost after insertion into ModuleList;
Runtime reflection confirms the underlying native object is indeed the target subclass (MLP), yet direct Java/Scala cast fails due to mismatched JVM wrapper class identities;
Since the returned reference is only typed as plain Module, we cannot invoke subclass-specific forward() methods or layer APIs without successful casting, breaking sequential forward propagation logic for stacked layers stored inside ModuleList.
Minimal Reproduction Code (Scala JavaCPP-PyTorch)
scala
import org.bytedeco.pytorch._
import torchrec.basic.layers.MLP

// 1. Create custom concrete layer subclass of Module
val mlpLayer = new MLP(inputDim=8, hiddenDims=Seq(16), outputDim=1)

// 2. Add to ModuleListImpl
val moduleList = new ModuleListImpl()
moduleList.push_back("mlp_block", mlpLayer)

// 3. Retrieve element, only typed as base Module
val retrieved: Module = moduleList.get(0)

// 4. ClassCastException thrown here
val casted: MLP = retrieved.asInstanceOf[MLP]
Expected Behavior
Elements fetched from ModuleListImpl should preserve runtime subclass type information of custom user-defined layers that extend Module;
Safe casting from the returned handle back to the original concrete layer subclass should be allowed without ClassCastException;
Provide a supported API to unwrap the actual custom layer implementation stored inside ModuleList, enabling access to subclass-specific forward() and layer attributes.
Workaround Limitation
Pure Java reflection cannot bypass this cast error, because the JNI-wrapped Module wrapper instance is a distinct JVM class from our custom Scala layer class, even if they point to the same underlying C++ module object. There is no trivial user-land workaround to restore the concrete layer type after retrieval from ModuleListImpl.

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

  // Use SequentialImpl like PyTorch's nn.Sequential
  private val sequential = new SequentialImpl()
  private var prevDim = inputDim
  private val layersRef = scala.collection.mutable.ListBuffer[AnyRef]()
  // Build MLP layers using SequentialImpl

  hiddenDims.foreach { dim =>
    sequential.push_back(new LinearImpl(prevDim, dim))

    if (useLayerNorm) {
      val vec = new LongVector(1)
      vec.put(0, dim)
      sequential.push_back(new LayerNormImpl(vec))
    } else if (useBatchNorm && !activation.equals("relu")) {
      sequential.push_back(new BatchNorm1dImpl(new BatchNormOptions(dim)))
    }
    activation.toLowerCase match {
      case "relu" => sequential.push_back(new ReLUImpl())
      case "sigmoid" =>sequential.push_back(new SigmoidImpl())
      case "tanh" => sequential.push_back(new TanhImpl())
      case "silu" | "swish" =>sequential.push_back(new SiLUImpl())
      case "gelu" => sequential.push_back(new GELUImpl())
      case "prelu" => sequential.push_back(new PReLUImpl())
      case "leaky_relu" | "leakyrelu" => sequential.push_back(new LeakyReLUImpl())
      case "none" | "identity" => sequential.push_back(new IdentityImpl())
      case _ =>sequential.push_back( new ReLUImpl())
    }

//    sequential.push_back(createActivation(activation))

    if (dropout > 0) {
      val dropoutLayer = new DropoutImpl(dropout)
      layersRef += dropoutLayer
      sequential.push_back(s"dropout_${layersRef.size}", dropoutLayer)
//      sequential.push_back(new DropoutImpl(dropout))
    }

    prevDim = dim
  }

  // Output layer
  if (outputLayer) {
    sequential.push_back(new LinearImpl(prevDim, outputDim))
  }

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

  def forward(x: Tensor): Tensor = {
    sequential.forward(x)
  }

  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()
    }
  }
}

class MetaLinear(
  inFeatures: Long,
  outFeatures: Long,
  device: String = DeviceSupport.backend
) extends Module {

  private val targetDevice = new Device(device)
  private val linear = new LinearImpl(inFeatures, outFeatures)
  linear.to(targetDevice, false)
  register_module("linear", linear)

  def forward(x: Tensor): Tensor = {
    linear.forward(x).to(targetDevice, ScalarType.Float)
  }

  def forwardFast(x: Tensor, fastWeight: Tensor, fastBias: Tensor): Tensor = {
    forward(x)
  }
}

  private val criticGates: ModuleListImpl = {
    val moduleList = new ModuleListImpl()
    for (i <- 0 until taskNum) {
      val gate = new MetaLinear(embedDim * 2, criticNum, device)
      gate.to(targetDevice, false)
      moduleList.push_back(gate)
    }
    moduleList
  }


   val criticGateOut = criticGates.get(taskIdx).asInstanceOf[MetaLinear].forward(criticGateInput)

--- MetaHeac Benchmark ---
Exception in thread "main" java.lang.ClassCastException: class org.bytedeco.pytorch.Module cannot be cast to class torchrec.basic.layers.MLP (org.bytedeco.pytorch.Module and torchrec.basic.layers.MLP are in unnamed module of loader 'app')
	at torchrec.models.multi_task.MetaHeac.$anonfun$7(MetaHeac.scala:186)
	at torchrec.models.multi_task.MetaHeac.$anonfun$adapted$3(MetaHeac.scala:185)
	at scala.collection.StrictOptimizedIterableOps.strictOptimizedMap(StrictOptimizedIterableOps.scala:102)
	at scala.collection.StrictOptimizedIterableOps.strictOptimizedMap$(StrictOptimizedIterableOps.scala:29)
	at scala.collection.immutable.Range.strictOptimizedMap(Range.scala:61)
	at scala.collection.StrictOptimizedIterableOps.map(StrictOptimizedIterableOps.scala:90)
	at scala.collection.StrictOptimizedIterableOps.map$(StrictOptimizedIterableOps.scala:29)
	at scala.collection.immutable.Range.map(Range.scala:213)
	at torchrec.models.multi_task.MetaHeac.forwardByName(MetaHeac.scala:185)
	at torchrec.trainers.MTLTrainer.fit$$anonfun$1(MTLTrainer.scala:71)
	at scala.runtime.java8.JFunction1$mcVI$sp.apply(JFunction1$mcVI$sp.scala:20)
	at scala.collection.immutable.Range.foreach(Range.scala:256)
	at torchrec.trainers.MTLTrainer.fit(MTLTrainer.scala:49)
	at benchmarks.BenchmarkRunner$.runMultiTaskBenchmark(BenchmarkRunner.scala:1307)
	at benchmarks.BenchmarkRunner$.runMetaHeacBenchmark(BenchmarkRunner.scala:754)
	at benchmarks.BenchmarkRunner$.main(BenchmarkRunner.scala:74)
	at benchmarks.BenchmarkRunner.main(BenchmarkRunner.scala)


```

[mullerhai]

Progress & Blockers on libtorch Module/ModuleDict Custom Binding Fixes
Labels: enhancement, libtorch, binding, cpp-customization

HI @saudet

I’d like to share the current progress, solved problems and remaining unresolved tasks when extending libtorch module APIs via JavaCPP custom C++ patch & binding logic.

1. Problems already fully resolved
2. Generic as() method implemented on base Module class
   We added a generic as() method on the top-level Module base type.
   This implementation solves the type casting pain point caused by Java generic type erasure:
   Any custom complex neural network model class that inherits from Module can be safely cast back to its original concrete subclass type, even if the object was stored/upcast to the base Module type in advance.
   No extra manual type conversion helper functions are required for downstream Scala/Java inference code.
3. Accessible ModuleDict::insert() mapping enabled via C++ source patch
   This issue does not belong to JavaCPP itself — the root cause is the original libtorch C++ source marks the insert() member function of ModuleDict as private, which makes it invisible to JavaCPP’s automatic binding parser.
   Current workaround solution:
   Modify the libtorch build script to forcibly rewrite the access modifier of ModuleDict::insert() from private to public during the pre-build C++ patching stage. After this patch, JavaCPP can successfully generate the JNI binding for ModuleDict.insert().
4. Unfinished / Blocked tasks (still under trial development)
   Two core extensions are still not working properly after multiple attempts, we are continuing debugging and hope to find a viable fix soon:
   Add a virtual forward() abstract method on the base Module class to unify model forward computation interfaces for all sub-modules.
   Implement three overloaded variants of push_back() inside SequentialImpl, which accept base Module type as input parameters for flexible sequential container construction.
5. Request / Expectation
   If anyone has experience customizing virtual method injection or overloaded container function binding in JavaCPP’s libtorch preset, kindly share related patch examples or configuration tips. We will update this issue once the two unfinished features are fully functional.

ModuleDict

  /** Updated the {@code ModuleDict} with key-value pairs from {@code OrderedDict} or
   *  {@code ModuleDict}. */
  /** Private {@code OrderedDict} holding the key-Module pairs. */
  public native @ByRef StringSharedModuleDict modules_(); public native ModuleDictImpl modules_(StringSharedModuleDict setter);

  /** Insert a key-module pair by overwriting existing keys,
   *  and register or replace the {@code Module}. */
  public native void insert(@StdString BytePointer key, @SharedPtr("torch::nn::Module") @ByVal Module module);
  public native void insert(@StdString String key, @SharedPtr("torch::nn::Module") @ByVal Module module);

Module

  /**
   * Attempts to cast this {@link Module} to any subclass passed via
   * {@code moduleClass}. For the 80+ built-in {@code *Impl} classes
   * this dispatches to the corresponding {@code asXxx()} method via
   * reflection so the C++ {@code dynamic_cast} actually verifies
   * the underlying object. For user-defined subclasses (custom
   * modules, complex wrapper classes that compose several layers,
   * …) the pointer is wrapped as-is: registration via
   * {@link #registerCppType(Class, String)} is optional but lets
   * the cast be type-checked when the user knows the C++ typeid of
   * their module.
   *
   * @param moduleClass non-null Java class extending {@link Module}.
   * @return a fresh instance of {@code moduleClass} wrapping this
   *         {@code Module}'s native pointer, or {@code null} when
   *         the underlying C++ object is not of the requested type.
   */
  public <M extends Module> M as(Class<M> moduleClass) {
      if (moduleClass == null) {
          throw new NullPointerException("moduleClass");
      }
      if (moduleClass.isInstance(this)) {
          return moduleClass.cast(this);
      }
      String className = moduleClass.getSimpleName();
      if (className.endsWith("Impl")) {
          String asMethodName = "as" + className.substring(0, className.length() - 4);
          try {
              java.lang.reflect.Method asMethod = Module.class.getMethod(asMethodName);
              Object module = asMethod.invoke(this);
              // Guard: the typed asXxx() returns a non-null wrapper only when
              // the underlying C++ object is actually of the requested type.
              // If moduleClass isn't a Java supertype of the returned wrapper,
              // surface that as a no-match instead of throwing ClassCastException.
              if (module == null) {
                  return null;
              }
              return moduleClass.cast(module);
          } catch (NoSuchMethodException e) {
              // Fall through to the generic Module pointer wrapper below.
          } catch (IllegalAccessException e) {
              throw new RuntimeException(e);
          } catch (java.lang.reflect.InvocationTargetException e) {
              Throwable cause = e.getCause();
              if (cause instanceof RuntimeException) {
                  throw (RuntimeException) cause;
              }
              throw new RuntimeException(cause);
          } catch (ClassCastException e) {
              // The asXxx() returned something that isn't a moduleClass
              // (e.g. user asked for FancyMLP but the underlying object
              // is a built-in layer). Treat as a no-match.
              return null;
          }
      }
      // Custom (non-*Impl) subclass: wrap the raw Module pointer.
      // The caller is responsible for knowing what they are doing;
      // if they want C++ typeid verification they should call
      // registerCppType() first.
      return ModuleAsHelper.wrap(this, moduleClass);
  }
  /**
   * Static registry mapping Java {@link Module} subclasses to fully
   * qualified C++ typeid names. Can be extended at runtime via
   * {@link #registerCppType(Class, String)} so custom user modules
   * can be type-checked at runtime as well.
   */
  private static final java.util.Map<Class<? extends Module>, String> CPP_TYPE_REGISTRY =
          new java.util.concurrent.ConcurrentHashMap<Class<? extends Module>, String>();
  /** Look up the C++ typeid registered for {@code javaClass}, if any. */
  public static String cppTypeName(Class<? extends Module> javaClass) {
      return CPP_TYPE_REGISTRY.get(javaClass);
  }
  /**
   * Register a custom Java {@link Module} subclass so that
   * {@link #as(Class)} can later look up the C++ typeid. Calling this
   * is optional — without it {@code as(Class)} will still return a
   * wrapper, just without typeid verification. {@code cppTypeName}
   * must be the fully qualified C++ type name, e.g.
   * {@code "torch::nn::LinearImpl"}.
   */
  public static <M extends Module> void registerCppType(Class<M> javaClass, String cppTypeName) {
      if (javaClass == null || cppTypeName == null) {
          throw new IllegalArgumentException("javaClass and cppTypeName must be non-null");
      }
      CPP_TYPE_REGISTRY.put(javaClass, cppTypeName);
  }


/**
 * Small reflection helper backing the generic
 * {@link Module#as(Class) Module.as(Class)} method.
 *
 * <p>The {@code as(Class)} method is generated by the {@code torch.java}
 * preset (see {@code torch::nn::Module::asModuleUnchecked}) and handles the
 * 80+ built-in {@code *Impl} classes through reflection on the existing
 * {@code asXxx()} family of methods. For <em>user-defined</em> subclasses
 * (custom modules, complex wrapper classes that compose several layers,
 * …) it falls back to {@link #wrap(Module, Class)}, which simply wraps the
 * existing native pointer with a fresh Java instance of the requested
 * class. Verification is up to the caller; {@link Module#registerCppType}
 * can be used to register the C++ typeid if the user wants the cast to be
 * validated through {@code dynamic_cast}.
 *
 * <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 directly — we deliberately do
     * <em>not</em> walk up to {@link Module}'s constructor because doing so
     * would silently produce a {@link Module} instance typed as the
     * requested subclass, which is a lie about the runtime type.
     */
    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);
        }
    }
}

[mullerhai]

now SequentialImpl could try support Module object push_back insert

  /**
   * Adds an existing {@link Module} (any subclass — built-in or
   * user-defined) to this {@code Sequential} container. The
   * underlying C++ type is identified at runtime via
   * {@code dynamic_cast} on the actual C++ object and the
   * matching typed {@code push_back<Impl>(name, ptr)} overload
   * is invoked. Throws {@link IllegalArgumentException} when the
   * underlying C++ type is not one of the recognised built-in
   * layers.
   */
  public native @Name("push_back") void push_back(@SharedPtr @Cast({"", "std::shared_ptr<torch::nn::Module>"}) Module module);
  public native @Name("push_back") void push_back(@StdString BytePointer name, @SharedPtr @Cast({"", "std::shared_ptr<torch::nn::Module>"}) Module module);
  public native @Name("push_back") void push_back(@StdString String name, @SharedPtr @Cast({"", "std::shared_ptr<torch::nn::Module>"}) Module module);

also will meet error

/Users/muller/Downloads/javacpp-presets-master/pytorch/cppbuild/macosx-arm64/include/torch/csrc/api/include/torch/nn/modules/container/any.h:219:7: error: static assertion failed due to requirement 'torch::detail::has_forward<torch::nn::Module>::value': Can only store module with a forward()
      method that has a non-templatized argument type and return type into AnyModule (e.g. we cannot store nn::Sequentialinto AnyModule, because its forward() method's argument type and return type are templatized. If you need to use nn::Sequentials inside each other you can subclass nn::Sequential and write a
      non-templatized forward function for it. You can checkout https://github.com/pytorch/vision/blob/2f46070f3cb1ea894d82578f3dc5677f82f34958/torchvision/csrc/models/mnasnet.cpp#L59 for an example on how to do this.).
  219 |       torch::detail::has_forward<ModuleType>::value,

/Users/muller/Downloads/javacpp-presets-master/pytorch/cppbuild/macosx-arm64/include/torch/csrc/api/include/torch/nn/modules/container/any.h:210:49: error: no member named 'forward' in 'torch::nn::Module'
  210 |           &std::remove_reference_t<ModuleType>::forward)) {