如何在Clojure中实现NDArray的加法运算符"+"？

问题场景 #

在Java中，给NDArray执行加法操作非常直接：

NDManager manager = NDManager.newBaseManager();
NDArray x = manager.create(3f);
NDArray y = manager.create(2f);
x.add(y);

但在Clojure中直接尝试类似操作时会报错：

(def ndm (nd/base-manager))
(def x (nd/create ndm 3.))
(def y (nd/create ndm 2.))

(nd/+ x y) ;; 报错：This NDArray implementation does not currently support this operation

报错原因是Clojure无法自动识别参数为NDArray类型，也不会自动映射到NDArray.add方法。

解决方案 #

你可以通过在自定义命名空间中重载+运算符，同时处理NDArray间相加、NDArray与标量相加，以及多参数相加的场景，完整实现如下：

(ns djl.ndarray
  (:import [ai.djl.ndarray NDArray])
  (:refer-clojure :exclude [+ - / *
                            = <= < >= >
                            identity to-array
                            min max concat
                            get set
                            flatten
                            sort]))

(defn +
  "实现NDArray加法，支持：
  - 单个NDArray直接返回
  - NDArray与NDArray相加
  - NDArray与标量（Number类型）相加
  - 多个参数链式相加"
  ([array] array)
  ([array0 array1]
   (cond
     (instance? NDArray array0)
     (.add array0 array1)
     (instance? NDArray array1)
     (.add array1 array0)
     :else
     (clojure.core/+ array0 array1)))
  ([array0 array1 & more]
   (reduce + (+ array0 array1) more)))

代码说明 #

1. 排除原生运算符：通过:refer-clojure :exclude排除Clojure原生的+，避免命名冲突
2. 多态处理：
   • 单参数场景直接返回输入，符合Clojure原生+的行为
   • 双参数时判断是否为NDArray，调用DJL提供的.add方法；非NDArray参数则回退到原生加法
   • 多参数通过reduce实现链式运算，支持任意数量的NDArray或标量相加
3. 兼容标量：DJL的NDArray.add本身支持接收Number类型的标量，因此无需额外转换即可直接使用

内容的提问来源于stack exchange，提问作者madeinQuant