fix: resolve all reflection warnings

src/clorch/amp.clj

@@ -4,13 +4,15 @@
   (:import [org.bytedeco.pytorch Optimizer Scalar Tensor TensorVector]
            [org.bytedeco.pytorch.global torch torch$DeviceType torch$ScalarType]))
 
-(defn- device-type [device]
+(set! *warn-on-reflection* true)
+
+(defn- device-type ^torch$DeviceType [device]
   (case device
     :cuda torch$DeviceType/CUDA
     :cpu torch$DeviceType/CPU
     (throw (ex-info "Autocast device must be :cuda or :cpu" {:device device}))))
 
-(defn- scalar-type [dtype]
+(defn- scalar-type ^torch$ScalarType [dtype]
   (case dtype
     :float16 torch$ScalarType/Half
     :bfloat16 torch$ScalarType/BFloat16

src/clorch/autograd.clj

@@ -2,6 +2,8 @@
   (:require [clorch.torch :as torch])
   (:import [org.bytedeco.pytorch NoGradGuard]))
 
+(set! *warn-on-reflection* true)
+
 (defn set-requires-grad [t bool]
   (.set_requires_grad (torch/->tensor t) (boolean bool))
   t)

src/clorch/cuda.clj

@@ -4,6 +4,8 @@
            [org.bytedeco.pytorch.global torch]
            [org.bytedeco.javacpp Loader]))
 
+(set! *warn-on-reflection* true)
+
 ;; Force loading with GPU extension if possible
 (let [old-ext (System/getProperty "org.bytedeco.javacpp.platform.extension")]
   (try

src/clorch/data.clj

@@ -6,6 +6,8 @@
            [java.util.concurrent Callable Future LinkedBlockingQueue ThreadFactory ThreadPoolExecutor TimeUnit TimeoutException]
            [java.util.concurrent.atomic AtomicInteger]))
 
+(set! *warn-on-reflection* true)
+
 (def ^:private torch-api
   (delay
     {:stack (requiring-resolve 'clorch.torch/stack)
@@ -38,8 +40,8 @@
     (let [cls (Class/forName "org.bytedeco.pytorch.JavaDataset")]
       (extend cls
         IDataset
-        {:get-size (fn [this] (.get (.size this)))
-         :get-item (fn [this idx]
+        {:get-size (fn [^org.bytedeco.pytorch.JavaDataset this] (.get (.size this)))
+         :get-item (fn [^org.bytedeco.pytorch.JavaDataset this idx]
                      (let [example (.get this (clojure.core/long idx))]
                        {:data (.data example) :target (.target example)}))}))))
 
@@ -90,7 +92,7 @@
       (instance? @tensor-class first-item)
       (let [v (org.bytedeco.pytorch.TensorVector.)]
         (doseq [item items] (.push_back v item))
-        (let [stacked ((torch-fn :stack) v 0)] (.retainReference stacked) stacked))
+        (let [stacked ((torch-fn :stack) v 0)] (.retainReference ^org.bytedeco.javacpp.Pointer stacked) stacked))
       :else (vec items))))
 
 (defn- build-batch [ds idxs collate-fn]
@@ -101,7 +103,7 @@
     (partition batch-size idxs)
     (partition-all batch-size idxs)))
 
-(defn- submit-batch [executor ds idxs collate-fn]
+(defn- submit-batch [^java.util.concurrent.ExecutorService executor ds idxs collate-fn]
   (.submit executor
            ^Callable
            (reify Callable
@@ -125,7 +127,7 @@
 
 (defn- parallel-batches
   [ds chunks collate-fn num-workers prefetch-factor]
-  (let [executor (make-thread-worker-executor num-workers)
+  (let [^ThreadPoolExecutor executor (make-thread-worker-executor num-workers)
         max-inflight (clojure.core/max 1 (* (clojure.core/int num-workers) (clojure.core/int prefetch-factor)))]
     (letfn [(fill-queue [remaining q]
               (loop [remaining remaining
@@ -150,7 +152,7 @@
     (.getAbsolutePath f)))
 
 (defn- process-start-worker [spec-file worker-id]
-  (let [proc-builder (ProcessBuilder. ["clojure" "-M" "-m" "clorch.data-worker" "--server" spec-file])
+  (let [proc-builder (ProcessBuilder. ^java.util.List (vector "clojure" "-M" "-m" "clorch.data-worker" "--server" spec-file))
         _ (.redirectError proc-builder java.lang.ProcessBuilder$Redirect/INHERIT)
         proc (.start proc-builder)
         rdr (BufferedReader. (InputStreamReader. (.getInputStream proc)))
@@ -280,7 +282,7 @@
         select (torch-fn :select)
         X-t (->tensor X)
         y-t (->tensor y)
-        n (.size X-t (clojure.core/long 0))]
+        n (.size ^org.bytedeco.pytorch.Tensor X-t (clojure.core/long 0))]
     (dataset :size (fn [] n)
              :get-item (fn [idx] {:data (select X-t 0 idx)
                                   :target (select y-t 0 idx)}))))

src/clorch/data_worker.clj

@@ -3,6 +3,8 @@
             [clojure.java.io :as io]
             [clorch.data :as data]))
 
+(set! *warn-on-reflection* true)
+
 (defn- load-dataset [spec-file]
   (let [{:keys [factory args worker-init-fn]} (edn/read-string (slurp spec-file))
         ctor (requiring-resolve factory)

src/clorch/dev.clj

@@ -5,6 +5,8 @@
   (:import [java.net BindException SocketException]
            [java.lang System]))
 
+(set! *warn-on-reflection* true)
+
 (defn- configure-javacpp-platform-extension! []
   (let [extension (platform/configure-platform-extension!)]
     (println "clorch.dev: org.bytedeco.javacpp.platform.extension=" (pr-str extension))))

src/clorch/distributed.clj

@@ -12,6 +12,8 @@
            [org.bytedeco.javacpp.chrono Milliseconds]
            [org.bytedeco.pytorch Store StringVector TCPStore TCPStoreOptions Tensor]))
 
+(set! *warn-on-reflection* true)
+
 (platform/configure-platform-extension!)
 
 (declare destroy-process-group!)
@@ -85,8 +87,8 @@
     resolved))
 
 (defn- tcp-store
-  [{:keys [rank world-size timeout-ms master-address master-port]}]
-  (let [timeout (Milliseconds. timeout-ms)
+  ^TCPStore [{:keys [rank world-size timeout-ms master-address master-port]}]
+  (let [timeout (Milliseconds. (long timeout-ms))
         workers (SizeTPointer. (long-array [world-size]))
         options (doto (TCPStoreOptions.)
                   (.port (unchecked-short master-port))
@@ -96,7 +98,7 @@
                   (.multiTenant false)
                   (.timeout timeout)
                   (.useLibUV true))]
-    (TCPStore. master-address options)))
+    (TCPStore. ^String master-address options)))
 
 (defn init-process-group!
   "Initializes this JVM's default NCCL process group.
@@ -313,7 +315,7 @@
     (.set store rank-key "ready")
     (dotimes [peer (:world-size context)]
       (.push_back store-keys (str prefix peer)))
-    (._wait store store-keys (Milliseconds. timeout-ms))))
+    (._wait store store-keys (Milliseconds. (long timeout-ms)))))
 
 (defn barrier!
   "Blocks until every rank enters the barrier."

src/clorch/distributed/checkpoint.clj

@@ -10,6 +10,8 @@
            [org.bytedeco.pytorch Generator]
            [org.bytedeco.pytorch.global torch]))
 
+(set! *warn-on-reflection* true)
+
 (def ^:private checkpoint-version 1)
 
 (defn- metadata-path [path]

src/clorch/distributed/launch.clj

@@ -10,6 +10,8 @@
            [java.util UUID]
            [java.util.concurrent LinkedBlockingQueue TimeUnit TimeoutException]))
 
+(set! *warn-on-reflection* true)
+
 (declare stop!)
 
 (defrecord WorkerProcess [rank ^Process process ^File stdout ^File stderr])
@@ -82,7 +84,7 @@
     (Files/createDirectories path (make-array java.nio.file.attribute.FileAttribute 0))
     (.toFile path)))
 
-(defn- write-worker-config! [config log-directory]
+(defn- write-worker-config! ^File [config log-directory]
   (let [path (.toPath (io/file log-directory "worker-config.edn"))
         content (ensure-edn! (select-keys config [:main :args]))]
     (Files/writeString path content StandardCharsets/UTF_8

src/clorch/distributed/nccl.clj

@@ -10,6 +10,8 @@
            [org.bytedeco.javacpp BytePointer IntPointer Pointer]
            [org.bytedeco.pytorch ByteVector Store TensorBase TensorVector]))
 
+(set! *warn-on-reflection* true)
+
 (defrecord NcclWork [^CUevent_st event ^ncclComm communicator result released?])
 
 (defn- nccl-message [status]
@@ -69,7 +71,7 @@
     (throw (ex-info "Unsupported NCCL reduction operation"
                     {:operation operation}))))
 
-(defn- unique-id-bytes [^ncclUniqueId unique-id]
+(defn- unique-id-bytes ^bytes [^ncclUniqueId unique-id]
   (byte-array
    (map #(.internal unique-id (int %))
         (range nccl/NCCL_UNIQUE_ID_BYTES))))
@@ -166,7 +168,7 @@
             (let [async-status (IntPointer. 1)]
               (check-nccl! "ncclCommGetAsyncError"
                            (nccl/ncclCommGetAsyncError
-                            (:communicator work) async-status))
+                            ^ncclComm (:communicator work) async-status))
               (check-nccl! "NCCL asynchronous operation" (.get async-status 0))
               (cudart/cudaEventDestroy (:event work))
               (reset! (:released? work) true))

src/clorch/distributed/worker.clj

@@ -4,6 +4,8 @@
             [clorch.cuda :as cuda]
             [clorch.distributed :as dist]))
 
+(set! *warn-on-reflection* true)
+
 (defn- ensure-java-25! []
   (let [feature (.feature (Runtime/version))]
     (when (< feature 25)

src/clorch/distributions.clj

@@ -8,8 +8,11 @@
   (:import [org.apache.commons.math3.distribution BetaDistribution ChiSquaredDistribution
             FDistribution GammaDistribution TDistribution WeibullDistribution ParetoDistribution]
            [org.apache.commons.math3.special Gamma]
+           [org.bytedeco.pytorch Tensor]
            [org.bytedeco.pytorch.global torch]))
 
+(set! *warn-on-reflection* true)
+
 (def ^:private eps 1e-8)
 (def ^:private log-2pi 1.8378770664093453)
 (def ^:private log-pi 1.1447298858494002)
@@ -43,7 +46,7 @@
 (defn- maybe-shape [sample-shape]
   (when (seq sample-shape) (vec sample-shape)))
 
-(defn- as-tensor [x]
+(defn- as-tensor ^Tensor [x]
   (if (number? x) (t/tensor x) (t/->tensor x)))
 
 (defn- const-like [value c]

src/clorch/einsum.clj

@@ -3,6 +3,8 @@
   (:require [clojure.string :as str]
             [clorch.torch :as torch]))
 
+(set! *warn-on-reflection* true)
+
 (def ^:private index-symbols
   "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
 

src/clorch/linalg.clj

@@ -6,6 +6,8 @@
   (:refer-clojure :exclude [cond])
   (:require [clorch.torch :as t]))
 
+(set! *warn-on-reflection* true)
+
 (defn norm [x] (t/linalg-norm x))
 (defn vector-norm [x] (t/linalg-vector-norm x))
 (defn matrix-norm [x] (t/linalg-matrix-norm x))