accelerating inference, let javacpp-pytorch support AOTIModelPackageLoader and AOTIModelContainerRunnerCuda and Profiler export file

[mullerhai]

HI， @saudet

For me, AOTIModelPackageLoader, AOT Cuda, and the profiler are critical factors for future inference acceleration. I am confident that we can use javacpp-pytorch to refactor a large model inference engine that rivals the performance of vllm, sglang, and deepspeed. A large model inference engine based on Java will one day be a tool on par with Apache Spark and Flink. However, if you think this is too ahead of its time, of course, if you consider that many Java programmers are still quite初级 (beginner), they are still stuck at the level of how to use tensor or matrix calculations. In any case, from the perspective of ecosystem completeness, AOTIModelPackageLoader, AOT Cuda, and the profiler are indispensable for javacpp-pytorch. If you can implement this, I am willing to pay 200 US dollars. This is not just for me; it is truly for the benefit of the Java, Scala, and Kotlin ecosystems. I believe that in the future, more major companies will also consider these aspects. With AOTIModelPackageLoader, AOT Cuda, and the profiler, I even feel that javacpp-pytorch would not need updates for the next five years, because it would truly be complete, without the need for other new, breakthrough features.
For everyone, these topics are already very profound, and many people may not even encounter them, or have never even considered them. However, from the perspective of accelerating inference, they are essential requirements.

[mullerhai]

hi saudet, we are really eager to have torch.tensor and torch.from_blob directly support multidimensional arrays as parameters. Flattening the arrays into one dimension is indeed not very convenient and also wastes more performance. It seems libtorch itself supports direct input of multidimensional arrays, so we kindly request that this feature be implemented.

import org.bytedeco.javacpp.*;
import org.bytedeco.pytorch.*;
import org.bytedeco.pytorch.global.torch;

import java.lang.reflect.Array;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

public class TensorConverter {

    public static void main(String[] args) {
        // 1. 测试 3D Float
        float[][][] floatData = {
                {{1.0f, 2.0f}, {3.0f, 4.0f}},
                {{5.0f, 6.0f}, {7.0f, 8.0f}}
        };
        Tensor tFloat = TensorConverter.toTensor(floatData);
        System.out.println("Float Tensor:");
        System.out.println(tFloat);
        // 输出应为 [2, 2, 2] 的 tensor

        // 2. 测试 2D Boolean (会自动转为 Bool 类型)
        boolean[][] boolData = {
                {true, false},
                {false, true}
        };
        Tensor tBool = TensorConverter.toTensor(boolData);
        System.out.println("Bool Tensor:");
        System.out.println(tBool);
        System.out.println("Dtype: " + tBool.dtype());
        // 输出应为 [2, 2]，dtype 为 Bool

        // 3. 测试 4D Int
        int[][][][] intData = new int[2][2][2][2];
        // ... 填充数据 ...
        Tensor tInt = TensorConverter.toTensor(intData);
        System.out.println("Int Tensor Shape: " + Arrays.toString(tInt.shape()));
    }
    /**
     * 通用入口：将任意维度的 Java 数组转换为 PyTorch Tensor
     * 支持类型：int, long, float, double, short, byte, boolean
     *
     * @param data 多维数组对象 (e.g. float[][], int[][][][])
     * @return 对应 shape 和 dtype 的 Tensor
     */
    public static Tensor toTensor(Object data) {
        if (data == null || !data.getClass().isArray()) {
            throw new IllegalArgumentException("Input data must be a Java array.");
        }

        // 1. 获取数组的 Shape 和总元素数量
        List<Long> shapeList = new ArrayList<>();
        getShape(data, shapeList);
        long[] shape = shapeList.stream().mapToLong(Long::longValue).toArray();
        int totalElements = 1;
        for (long s : shape) {
            totalElements *= (int) s;
        }

        // 2. 获取底层元素类型
        Class<?> componentType = getRootComponentType(data.getClass());

        // 3. 根据类型进行分发处理
        if (componentType == float.class) {
            float[] flat = new float[totalElements];
            flattenFloat(data, flat, new int[]{0});
            return torch.tensor(new FloatPointer(flat)).view(shape);
        }
        else if (componentType == int.class) {
            int[] flat = new int[totalElements];
            flattenInt(data, flat, new int[]{0});
            return torch.tensor(new IntPointer(flat)).view(shape);
        }
        else if (componentType == double.class) {
            double[] flat = new double[totalElements];
            flattenDouble(data, flat, new int[]{0});
            return torch.tensor(new DoublePointer(flat)).view(shape);
        }
        else if (componentType == long.class) {
            long[] flat = new long[totalElements];
            flattenLong(data, flat, new int[]{0});
            return torch.tensor(new LongPointer(flat)).view(shape);
        }
        else if (componentType == byte.class) {
            byte[] flat = new byte[totalElements];
            flattenByte(data, flat, new int[]{0});
            return torch.tensor(new BytePointer(flat)).view(shape);
        }
        else if (componentType == short.class) {
            short[] flat = new short[totalElements];
            flattenShort(data, flat, new int[]{0});
            return torch.tensor(new ShortPointer(flat)).view(shape);
        }
        else if (componentType == boolean.class) {
            // 特殊处理 Boolean：转换为 Byte (0/1) 并指定 kBool
            byte[] flat = new byte[totalElements];
            flattenBoolean(data, flat, new int[]{0});

            // 关键：创建 Tensor 时指定 dtype 为 kBool
            TensorOptions options = new TensorOptions().dtype(new ScalarTypeOptional(torch.ScalarType.Float));
            // 注意：这里先通过 BytePointer 创建，然后告知 PyTorch 它是 Bool
            // JavaCPP 的 torch.tensor(BytePointer) 默认创建 kByte (uint8/int8)
            // 所以我们需要显式使用工厂方法或 cast

            // 方法 A: 先生成 Byte Tensor，再 cast (推荐，比较稳妥)
            Tensor byteTensor = torch.tensor(new BytePointer(flat));
            return byteTensor.to(torch.ScalarType.Bool).view(shape);

            // 方法 B: 直接使用 blob (需要小心内存生命周期，这里为了安全用方法 A 拷贝一次)
        }
        else if (componentType == String.class) {
            throw new UnsupportedOperationException(
                    "Direct String array to Tensor conversion is not supported in LibTorch C++ API. " +
                            "Please tokenize your strings into integers (longs) or encode them as UTF-8 bytes first."
            );
        }
        else {
            throw new IllegalArgumentException("Unsupported data type: " + componentType.getName());
        }
    }

    // ==========================================================
    //                 私有辅助方法：递归展平
    //   使用 int[] index 包装器来在递归中保持计数，避免 ArrayList 装箱
    // ==========================================================

    private static void flattenFloat(Object array, float[] flat, int[] index) {
        if (array instanceof float[]) {
            float[] arr = (float[]) array;
            System.arraycopy(arr, 0, flat, index[0], arr.length);
            index[0] += arr.length;
        } else {
            Object[] objArr = (Object[]) array;
            for (Object o : objArr) flattenFloat(o, flat, index);
        }
    }

    private static void flattenInt(Object array, int[] flat, int[] index) {
        if (array instanceof int[]) {
            int[] arr = (int[]) array;
            System.arraycopy(arr, 0, flat, index[0], arr.length);
            index[0] += arr.length;
        } else {
            Object[] objArr = (Object[]) array;
            for (Object o : objArr) flattenInt(o, flat, index);
        }
    }

    private static void flattenDouble(Object array, double[] flat, int[] index) {
        if (array instanceof double[]) {
            double[] arr = (double[]) array;
            System.arraycopy(arr, 0, flat, index[0], arr.length);
            index[0] += arr.length;
        } else {
            Object[] objArr = (Object[]) array;
            for (Object o : objArr) flattenDouble(o, flat, index);
        }
    }

    private static void flattenLong(Object array, long[] flat, int[] index) {
        if (array instanceof long[]) {
            long[] arr = (long[]) array;
            System.arraycopy(arr, 0, flat, index[0], arr.length);
            index[0] += arr.length;
        } else {
            Object[] objArr = (Object[]) array;
            for (Object o : objArr) flattenLong(o, flat, index);
        }
    }

    private static void flattenShort(Object array, short[] flat, int[] index) {
        if (array instanceof short[]) {
            short[] arr = (short[]) array;
            System.arraycopy(arr, 0, flat, index[0], arr.length);
            index[0] += arr.length;
        } else {
            Object[] objArr = (Object[]) array;
            for (Object o : objArr) flattenShort(o, flat, index);
        }
    }

    private static void flattenByte(Object array, byte[] flat, int[] index) {
        if (array instanceof byte[]) {
            byte[] arr = (byte[]) array;
            System.arraycopy(arr, 0, flat, index[0], arr.length);
            index[0] += arr.length;
        } else {
            Object[] objArr = (Object[]) array;
            for (Object o : objArr) flattenByte(o, flat, index);
        }
    }

    private static void flattenBoolean(Object array, byte[] flat, int[] index) {
        if (array instanceof boolean[]) {
            boolean[] arr = (boolean[]) array;
            for (boolean b : arr) {
                flat[index[0]++] = b ? (byte) 1 : (byte) 0;
            }
        } else {
            Object[] objArr = (Object[]) array;
            for (Object o : objArr) flattenBoolean(o, flat, index);
        }
    }

    // ==========================================================
    //                 私有辅助方法：Shape 计算
    // ==========================================================

    private static void getShape(Object array, List<Long> shape) {
        if (array.getClass().isArray()) {
            shape.add((long) Array.getLength(array));
            if (Array.getLength(array) > 0) {
                // 递归检查第一个元素来确定下一维的长度
                // 假设数组是规则的（即每一行的长度相同）
                getShape(Array.get(array, 0), shape);
            }
        }
    }

    private static Class<?> getRootComponentType(Class<?> arrayClass) {
        Class<?> type = arrayClass;
        while (type.isArray()) {
            type = type.getComponentType();
        }
        return type;
    }
}

[mullerhai]

in cpp libtorch

	auto B = torch::tensor({{1, 2, 3}, {2, 0, 4}, {3, 4, 5}});
	std::cout << "B = \n" << B << std::endl;

invoke chain

inline at::Tensor tensor(detail::TensorDataContainer tensor_data_container, const at::TensorOptions& options = {}) {
  return autograd::make_variable(
    // note: we remove the requires_grad setting from the TensorOptions because
    // it is ignored anyways (and we actually have an assertion that it isn't set
    // which would fail otherwise). We handle requires_grad explicitly here
    // instead of passing it through to the kernel.
    tensor_data_container.convert_to_tensor(options.requires_grad(::std::nullopt)),
    options.requires_grad());
}


#pragma once

#include <ATen/Dispatch.h>
#include <ATen/ScalarOps.h>
#include <ATen/core/Tensor.h>
#include <ATen/core/grad_mode.h>

#include <c10/util/irange.h>

#ifndef AT_PER_OPERATOR_HEADERS
#include <ATen/Functions.h>
#else
#include <ATen/ops/empty.h>
#include <ATen/ops/tensor.h>
#endif

#include <initializer_list>

namespace torch::detail {

enum class TensorDataContainerType { Scalar, InitList, Tensor };

struct TensorDataContainer;

inline std::ostream& operator<<(
    std::ostream& stream,
    const TensorDataContainer& tensor_data_container);

inline c10::ScalarType compute_desired_dtype(c10::ScalarType scalar_type) {
  if (scalar_type == at::kInt || scalar_type == at::kLong) {
    // C++ `torch::tensor` with an integer type or an `at::ArrayRef` /
    // `std::vector` / (nested) braced-init-list of integer types always
    // produces a tensor of dtype `at::kLong` (aka. int64_t), matching Python
    // `torch.tensor` behavior.
    return at::kLong;
  } else if (scalar_type == at::kFloat || scalar_type == at::kDouble) {
    // C++ `torch::tensor` with a floating-point type or an `at::ArrayRef` /
    // `std::vector` / (nested) braced-init-list of floating-point types always
    // produces a tensor of dtype `torch::get_default_dtype()`, matching Python
    // `torch.tensor` behavior.
    return at::typeMetaToScalarType(at::get_default_dtype());
  } else {
    return scalar_type;
  }
}

// We use `TensorDataContainer` to support converting the following data
// container types into the equivalent Tensor:
//
// 1. Arbitrarily nested braced-init-list (e.g. `{{1, 2}, {3, 4}}`).
// 2. `at::ArrayRef` of supported tensor data types.
// 3. `std::vector` of supported tensor data types.
//
// At any time, a `TensorDataContainer` object represents one of the following:
//
// 1. A scalar with value `scalar()` and type `scalar_type()`.
// 2. A Tensor represented in `std::initializer_list<TensorDataContainer>` form,
//    with value `init_list()`, Tensor scalar type `scalar_type()`, and Tensor
//    sizes `sizes()`.
// 3. A Tensor represented in `at::Tensor` form, with value `tensor()`, scalar
// type `scalar_type()`,
//    and Tensor sizes `sizes()`.
//
// All the infrastructure here is mostly to support converting an arbitrarily
// nested braced-init-list to the equivalent Tensor successfully. Consider the
// following example:
//
// `torch::tensor({{1}, {2}})`
//
// this will call into the `torch::tensor` function:
//
// `at::Tensor tensor(detail::TensorDataContainer tensor_data_container, const
// at::TensorOptions& options = {})`
//
// the compiler will first try to convert `{{1}, {2}}` to `TensorDataContainer`
// type:
//
// `TensorDataContainer({{1}, {2}})`
//
// which matches to the
// `TensorDataContainer(std::initializer_list<TensorDataContainer>)`
// constructor, and in an attempt to convert `{1}` and `{2}` to
// `TensorDataContainer`, it calls the following:
//
// `TensorDataContainer({1})`  (same call path happens for `{2}`, and we'll just
// focus on `{1}` here)
//
// At this point, theoretically there are two plausible ways for `{1}` to be
// matched to one of the constructors of `TensorDataContainer`:
//
// 1. It can be a list-initialization of a scalar value, thus matching
// `TensorDataContainer(int value)`.
// 2. It can be converted to `std::initializer_list<TensorDataContainer>`, thus
// matching
//    `TensorDataContainer(std::initializer_list<TensorDataContainer>)`.
//
// How does the compiler decide which one to choose? According to
// `https://en.cppreference.com/w/cpp/language/list_initialization`,
// braced-init-list always prefers the constructor that takes
// `std::initializer_list`. Hence we happily move forward with constructor #2,
// and it calls the following:
//
// `TensorDataContainer(1)`
//
// Now it matches `TensorDataContainer(int value)`, which stores `1` as a scalar
// value. All is good.
struct TensorDataContainer {
  // NOTE: For tensors with zero-size dimensions (e.g. `torch::tensor({{},
  // {}})`), the innermost empty braced-init-list `{}` matches the default
  // constructor of the innermost `TensorDataContainer`.
  TensorDataContainer()
      : sizes_({0}),
        // NOTE: In Python, the dtype of tensors with zero-size dimensions (e.g.
        // `torch.tensor([[], []])`) depends on the value of
        // `torch.get_default_dtype()`, and we should do the same for the C++
        // equivalent.
        scalar_type_(at::typeMetaToScalarType(at::get_default_dtype())),
        type_(TensorDataContainerType::InitList) {}
#define TENSOR(T, S)                            \
  TensorDataContainer(T value)                  \
      : scalar_type_(at::k##S),                 \
        type_(TensorDataContainerType::Scalar), \
        scalar_(value) {}
  AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, TENSOR)
  AT_FORALL_COMPLEX_TYPES(TENSOR)
#undef TENSOR
  TensorDataContainer(std::initializer_list<TensorDataContainer> init_list)
      : scalar_type_(init_list.begin()->scalar_type()),
        type_(TensorDataContainerType::InitList),
        init_list_(init_list) {
    const TensorDataContainer& first_elem = *(init_list.begin());
    for (const auto& elem : init_list) {
      TORCH_CHECK(
          elem.sizes() == first_elem.sizes(),
          "Expected all sub-lists to have sizes: ",
          first_elem.sizes(),
          " (e.g. ",
          first_elem,
          "), ",
          "but got sub-list ",
          elem,
          " with sizes: ",
          elem.sizes());
      TORCH_CHECK(
          elem.scalar_type() == first_elem.scalar_type(),
          "Expected all elements of the tensor to have the same scalar type: ",
          first_elem.scalar_type(),
          ", but got element of scalar type: ",
          elem.scalar_type());
    }
    sizes_.reserve(first_elem.sizes().size() + 1);
    sizes_.push_back(static_cast<int64_t>(init_list.size()));
    sizes_.insert(
        sizes_.end(), first_elem.sizes().begin(), first_elem.sizes().end());
  }

#define TENSOR(T, S)                                                          \
  TensorDataContainer(at::ArrayRef<T> values)                                 \
      : sizes_({(int64_t)values.size()}),                                     \
        scalar_type_(at::k##S),                                               \
        type_(TensorDataContainerType::Tensor) {                              \
    at::AutoDispatchBelowAutograd mode;                                       \
    if (scalar_type_ == at::kBool) {                                          \
      tensor_ = at::tensor(values, at::TensorOptions().device(at::kCPU));     \
    } else {                                                                  \
      tensor_ = at::tensor(values, at::dtype(scalar_type_).device(at::kCPU)); \
    }                                                                         \
  }
  AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, TENSOR)
  AT_FORALL_COMPLEX_TYPES(TENSOR)
#undef TENSOR

  // NOTE: We need to handle `std::vector` explicitly instead of relying on an
  // implicit conversion to `at::ArrayRef`, otherwise the following error can be
  // thrown when calling `torch::tensor(std::vector<int>({1, 2}))`:
  // ```
  // error: no matching function for call to 'tensor(const std::vector<int>&)'
  // no known conversion for argument 1 from 'const std::vector<int>' to
  // 'torch::detail::TensorDataContainer'
  // ```
  //
  // NOTE: `torch::tensor(std::vector<bool>)` is not supported for now, because
  // ArrayRef<bool> cannot be constructed from a std::vector<bool> bitfield.
#define TENSOR(T, S)                                \
  TensorDataContainer(const std::vector<T>& values) \
      : TensorDataContainer(at::ArrayRef<T>(values)) {}
  AT_FORALL_SCALAR_TYPES_AND2(Half, BFloat16, TENSOR)
  AT_FORALL_COMPLEX_TYPES(TENSOR)
#undef TENSOR

  bool is_scalar() const {
    return type_ == TensorDataContainerType::Scalar;
  }

  const c10::Scalar& scalar() const {
    TORCH_CHECK(
        is_scalar(),
        "Can only call `scalar()` on a TensorDataContainer that has `is_scalar() == true`");
    return scalar_;
  }

  bool is_init_list() const {
    return type_ == TensorDataContainerType::InitList;
  }

  const std::initializer_list<TensorDataContainer>& init_list() const {
    TORCH_CHECK(
        is_init_list(),
        "Can only call `init_list()` on a TensorDataContainer that has `is_init_list() == true`");
    return init_list_;
  }

  bool is_tensor() const {
    return type_ == TensorDataContainerType::Tensor;
  }

  const at::Tensor& tensor() const {
    TORCH_CHECK(
        is_tensor(),
        "Can only call `tensor()` on a TensorDataContainer that has `is_tensor() == true`");
    return tensor_;
  }

  const std::vector<int64_t>& sizes() const {
    return sizes_;
  }

  const c10::ScalarType& scalar_type() const {
    return scalar_type_;
  }

  at::Tensor convert_to_tensor(at::TensorOptions options) const {
    if (!options.has_dtype()) {
      options = options.dtype(compute_desired_dtype(scalar_type_));
    }

    if (is_scalar()) {
      at::AutoDispatchBelowAutograd mode;
      return at::scalar_tensor(scalar_, options);
    } else if (is_init_list()) {
      // NOTE: Here we explicitly choose to initialize the tensor on CPU first,
      // fill each element of the tensor, and then move the tensor to the
      // desired device. For CUDA device, this approach only involves 1 CUDA
      // kernel launch, and is much faster than initializing the tensor on CUDA
      // first and then filling each element of it (which involves `N` CUDA
      // kernel launches where `N` is the number of the elements in the tensor).
      at::Tensor tensor = ([&]() {
        at::AutoDispatchBelowAutograd mode;
        return at::empty(sizes_, options.device(at::kCPU));
      })();
      fill_tensor(tensor);
      return tensor.to(options.device());
    } else if (is_tensor()) {
      auto output = tensor_.to(options);
      TORCH_CHECK(
          !tensor_.is_complex() || output.is_complex(),
          "can not do torch::tensor(complex, dtype=non-complex) because complex can not be casted to real number without loss of information");
      return output;
    } else {
      TORCH_INTERNAL_ASSERT(false, "Invalid TensorDataContainer type");
    }
  }

  void pretty_print_recursive(std::ostream& stream) const {
    if (is_scalar()) {
      AT_DISPATCH_ALL_TYPES_AND3(
          at::kBool,
          at::kHalf,
          at::kBFloat16,
          scalar_type_,
          "TensorDataContainer_pretty_print_scalar",
          [&] { stream << scalar_.to<scalar_t>(); });
    } else if (is_init_list()) {
      stream << "{";
      for (const TensorDataContainer* it = init_list_.begin();
           it != init_list_.end();
           it++) {
        stream << *it;
        if (std::next(it) != init_list_.end())
          stream << ", ";
      }
      stream << "}";
    } else if (is_tensor()) {
      stream << "{";
      for (const auto i : c10::irange(tensor_.sizes()[0])) {
        AT_DISPATCH_ALL_TYPES_AND3(
            at::kBool,
            at::kHalf,
            at::kBFloat16,
            scalar_type_,
            "TensorDataContainer_pretty_print_tensor_item",
            [&] { stream << tensor_[i].item<scalar_t>(); });
        if (i != tensor_.sizes()[0] - 1)
          stream << ", ";
      }
      stream << "}";
    } else {
      TORCH_INTERNAL_ASSERT(false, "Invalid TensorDataContainer type");
    }
  }

 private:
  void fill_tensor(at::Tensor& tensor) const {
    if (is_scalar()) {
      TORCH_INTERNAL_ASSERT(
          tensor.dim() == 0,
          "Expected a 0-dim Tensor, but got Tensor with dimensions: ",
          tensor.dim());
      at::NoGradGuard guard;
      tensor.fill_(scalar_);
    } else if (is_init_list()) {
      TORCH_INTERNAL_ASSERT(
          tensor.sizes()[0] == (int64_t)init_list_.size(),
          "Expected a Tensor with size ",
          init_list_.size(),
          " in its first dimension, but got Tensor with size ",
          tensor.sizes()[0],
          " in its first dimension");
      int64_t index = 0;
      for (const auto& elem : init_list_) {
        at::Tensor slice = tensor[index];
        elem.fill_tensor(slice);
        index++;
      }
    } else if (is_tensor()) {
      TORCH_INTERNAL_ASSERT(
          false,
          "TensorDataContainer is already a Tensor type, `fill_tensor` should not be called");
    } else {
      TORCH_INTERNAL_ASSERT(false, "Invalid TensorDataContainer type");
    }
  }

  std::vector<int64_t> sizes_;
  c10::ScalarType scalar_type_;
  TensorDataContainerType type_;
  c10::Scalar scalar_;
  std::initializer_list<TensorDataContainer> init_list_;
  at::Tensor tensor_;
};

inline std::ostream& operator<<(
    std::ostream& stream,
    const TensorDataContainer& tensor_data_container) {
  tensor_data_container.pretty_print_recursive(stream);
  return stream;
}

} // namespace torch::detail

in javacpp-pytorch

@Namespace("torch") public static native @ByVal Tensor tensor(@ByVal @Cast("torch::detail::TensorDataContainer*") Pointer tensor_data_container, @Const @ByRef(nullValue = "at::TensorOptions{}") TensorOptions options);
@Namespace("torch") public static native @ByVal Tensor tensor(@ByVal @Cast("torch::detail::TensorDataContainer*") Pointer tensor_data_container);

so. I want to. do we need implement TensorDataContainer to make Multidimensional array as. input parameter for torch.tensor() ? or make pointer could parse Multidimensional array ? @saudet

[mullerhai]

I write a. tool to get the shape and flatten the array , if necessay ,you could. add in

import java.lang.reflect.Array;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

public class TensorToolkit {

    // Get Multidimensional array Shape
    public static long[] getShape(Object array) {
        if (array == null || !array.getClass().isArray()) return new long[0];
        List<Long> shape = new ArrayList<>();
        Object current = array;
        while (current != null && current.getClass().isArray()) {
            int len = Array.getLength(current);
            shape.add((long) len);
            current = (len > 0) ? Array.get(current, 0) : null;
        }
        return shape.stream().mapToLong(i -> i).toArray();
    }

    // flatten Multidimensional array
    public static Object flatten(Object multiDimArray) {
        if (multiDimArray == null) return null;
        Class<?> componentType = multiDimArray.getClass();
        while (componentType.isArray()) componentType = componentType.getComponentType();

        int size = 1;
        for (long d : getShape(multiDimArray)) size *= d;

        Object result = Array.newInstance(componentType, size);
        fill(multiDimArray, result, new int[]{0});
        return result;
    }

    private static void fill(Object src, Object dest, int[] idx) {
        int len = Array.getLength(src);
        for (int i = 0; i < len; i++) {
            Object el = Array.get(src, i);
            if (el != null && el.getClass().isArray()) fill(el, dest, idx);
            else Array.set(dest, idx[0]++, el);
        }
    }
}

test

package com.example.neo.mcp.test.service;

import java.util.Arrays;

public class MultiDimTest {
    public static void main(String[] args) {

        // --- 6维示例 (int): 2x1x1x1x1x2 (共4个元素) ---
        // 结构: 2个 [1][1][1][1][2]
        int[][][][][][] arr6d = {
                {{{{{1, 2}}}}},
                {{{{{3, 4}}}}}
        };
        System.out.println("6D Shape: " + Arrays.toString(TensorToolkit.getShape(arr6d)));
        System.out.println("6D Flat:  " + Arrays.toString((int[]) TensorToolkit.flatten(arr6d)));
        System.out.println("---");

        // --- 7维示例 (short): 1x1x1x1x1x2x2 (共4个元素) ---
        // 结构: 1个 [1][1][1][1][2][2]
        short[][][][][][][] arr7d = {
                {{{{{{10, 20}, {30, 40}}}}}}
        };
        System.out.println("7D Shape: " + Arrays.toString(TensorToolkit.getShape(arr7d)));
        System.out.println("7D Flat:  " + Arrays.toString((short[]) TensorToolkit.flatten(arr7d)));
        System.out.println("---");

        // --- 8维示例 (float): 1x1x1x1x1x1x2x2 (共4个元素) ---
        float[][][][][][][][] arr8d = {
                {{{{{{{1.1f, 1.2f}, {1.3f, 1.4f}}}}}}}
        };
        System.out.println("8D Shape: " + Arrays.toString(TensorToolkit.getShape(arr8d)));
        System.out.println("8D Flat:  " + Arrays.toString((float[]) TensorToolkit.flatten(arr8d)));
        System.out.println("---");

        // --- 9维示例 (double): 1x1x1x1x1x1x1x2x2 (共4个元素) ---
        double[][][][][][][][][] arr9d = {
                {{{{{{{{0.1, 0.2}, {0.3, 0.4}}}}}}}}
        };
        System.out.println("9D Shape: " + Arrays.toString(TensorToolkit.getShape(arr9d)));
        System.out.println("9D Flat:  " + Arrays.toString((double[]) TensorToolkit.flatten(arr9d)));
        System.out.println("---");

        // --- 10维示例 (long): 2x1x1x1x1x1x1x1x1x2 (共4个元素) ---
        // 为了安全，10维建议直接用 new 初始化，然后手动赋值
        long[][][][][][][][][][] arr10d = new long[2][1][1][1][1][1][1][1][1][2];
        arr10d[0][0][0][0][0][0][0][0][0][0] = 100L;
        arr10d[0][0][0][0][0][0][0][0][0][1] = 200L;
        arr10d[1][0][0][0][0][0][0][0][0][0] = 300L;
        arr10d[1][0][0][0][0][0][0][0][0][1] = 400L;

        System.out.println("10D Shape: " + Arrays.toString(TensorToolkit.getShape(arr10d)));
        System.out.println("10D Flat:  " + Arrays.toString((long[]) TensorToolkit.flatten(arr10d)));
    }
}

console

6D Shape: [2, 1, 1, 1, 1, 2]
6D Flat:  [1, 2, 3, 4]
---
7D Shape: [1, 1, 1, 1, 1, 2, 2]
7D Flat:  [10, 20, 30, 40]
---
8D Shape: [1, 1, 1, 1, 1, 1, 2, 2]
8D Flat:  [1.1, 1.2, 1.3, 1.4]
---
9D Shape: [1, 1, 1, 1, 1, 1, 1, 2, 2]
9D Flat:  [0.1, 0.2, 0.3, 0.4]
---
10D Shape: [2, 1, 1, 1, 1, 1, 1, 1, 1, 2]
10D Flat:  [100, 200, 300, 400]
``

[saudet]

Java SE 8 comes with functionality for this, so I'd just use that:

import java.util.Arrays;
import java.util.stream.IntStream;

int[][] nestedArray = {{1, 2, 3}, {4, 5}, {6, 7, 8, 9}};

int[] flatArray = Arrays.stream(nestedArray) // Stream<int[]>
                        .flatMapToInt(Arrays::stream)   // IntStream
                        .toArray(); // int[]

// Output: [1, 2, 3, 4, 5, 6, 7, 8, 9]
System.out.println(Arrays.toString(flatArray));