Skip to content

Split tensor ops into functions - #168

Merged
benikm91 merged 2 commits into
dimwit-dev:mainfrom
benikm91:make-functions-for-extensions
Sep 27, 2026
Merged

benikm91 merged 2 commits into
dimwit-dev:mainfrom
benikm91:make-functions-for-extensions

Conversation

@benikm91

@benikm91 benikm91 commented Sep 26, 2026 •

Copy link
Copy Markdown
Collaborator

Split tensor ops into functions (*Ops) and extension methods (*Extensions)

Until now most tensor operations existed only as extension methods (t.relu, t.sum(Axis[A])), with no function you could pass around (e.g. to TreeOf.map). This PR adds a function for every operation and restructures tensorops consistently.

Structure

  • Every tensorops/X.scala now has XOps (the functions) and XExtensions (extension methods that only call XOps). They can't share one object: a function and an extension method with the same name have the same erased signature.
  • Functions are exported on the companion objects, e.g. Tensor.relu(t), Tensor.sum(t, Axis[A]), Tensor.softmax(t, Axis[A]), Tensor.dot(t1, Axis[A], t2), Tensor.conv2d(input, kernel). The Tensor1 => Tensor1 functions stay on Tensor1 (Tensor1.softmax). Users access operations as functions on Tensor/Tensor1 (e.g. Tensor.relu(t)) or as extension methods via import dimwit.* (e.g. t.relu). XOps and XExtensions are private[dimwit] and only structure the code internally.
  • Functions without a same-named extension or operator stay top-level via import dimwit.*: maximum, minimum, maximum_!, minimum_!, where, where_!, triu, tril, stack, concatenate, zipvmap.
  • TensorOps renamed to ValueTypeClasses: it only holds IsFloating, IsNumber, … now. TensorOpsUtil renamed to dimwit.tensor.Broadcast (public, top-level). Padding and Stride1/2/3 moved to dimwit.tensor. ValueOps renamed to ValueExtensions.

New

  • Broadcasting functions: add_!, subtract_!, multiply_!, divide_!, mod_!, less_!, … equal_!, logicalAnd_!/Or_!/Xor_!, maximum_!, minimum_!, arrayEqual_!. The ! operators call them.
  • ===!, approxEquals_!, approxElementEquals_!.
  • Scalar-first operators mirror scalar-last exactly, for all Scala scalar types: 3 %! t and 2.0 *! t now work, and the scalar takes the tensor's precision.

Fixes

  • Non-! two-tensor ops (+, <, maximum, where, ===, …) now fail fast when the extents differ. Before, JAX silently broadcast axes of extent 1 (e.g. 2×2 + 1×2). This didn't align with DimWit's intent of explicit broadcasting and no extent broadcasting.
  • approxElementEquals really compares elementwise (jnp.isclose). Before, it returned a single boolean (allclose) typed as a tensor.
  • Removed argsort(axes: Tuple). It always failed at runtime, because jnp.argsort only takes one axis.

Breaking changes

  • add, subtract, multiply, divide, mod, negate, less, lessEqual, greater, greaterEqual, equal and logicalAnd/Or/Xor/Not are no longer top-level: use the operators, or Tensor.add(...) etc.
  • Removed addScalar, subtractScalar, divideScalar, modScalar: use +!, -!, /!, %! with a Tensor0. multiplyScalar → Tensor.scale.
  • Scalar-first operators (2.0f *! t) now need import dimwit.Conversions.given, like t *! 2.0f already did.
  • quantile(q) takes a Tensor0[V] instead of a Float.
  • Removed retag (unused).
  • Direct imports of internals break: dimwit.tensor.TensorOps.*, dimwit.tensor.tensorops.*, ValueOps. Code using import dimwit.* only sees the changes above.

@benikm91

benikm91 commented Sep 26, 2026 •

Copy link
Copy Markdown
Collaborator Author

Note: Coding was done by Claude; however, it took several iterations, with me reviewing/git diff-ing the code. I will review this once more in detail. I wanted to have it as a PR already so I could share it and maybe get initial discussion points if needed.

@marcelluethi, so maybe go over the PR description, not the code, for now. The changes are extensive but mostly internal.

@benikm91
benikm91 marked this pull request as draft September 26, 2026 10:02

@marcelluethi marcelluethi left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I haven´t carefully checked all code but went through the most crucial places, like Tensor TensorNExtension, ValueExtension, AlongAxisOps. The new structure makes a lot of sense and I actually like that all functions on Tensor are now only available using the Tensor object and not globally. This should make both the internal organisation as well as the user experience better.

@benikm91
benikm91 marked this pull request as ready for review September 27, 2026 06:16
…ions)

Every tensor operation now also exists as a function, e.g. Tensor.relu(t),
Tensor.sum(t, Axis[A]), Tensor.softmax(t, Axis[A]), Tensor.dot(t1, Axis[A], t2).
Each tensorops file has an XOps object (functions, exported on Tensor/Tensor1)
and an XExtensions object (extension methods calling XOps, exported from
package.scala). All tensorops objects are private[dimwit].

- TensorOps -> ValueTypeClasses, TensorOpsUtil -> dimwit.tensor.Broadcast,
  ValueOps -> ValueExtensions; Padding/Stride moved to dimwit.tensor
- Add broadcasting functions (add_!, less_!, maximum_!, arrayEqual_!, ...),
  ===!, approxEquals_! and approxElementEquals_!
- Scalar-first operators mirror scalar-last for all Scala scalar types
- Non-! binary ops fail fast on differing extents instead of JAX
  silently broadcasting axes of extent 1
- Fix approxElementEquals returning a single boolean (allclose -> isclose)
- Remove broken argsort(axes: Tuple), unused retag and the *Scalar
  functions; quantile takes a Tensor0; add/less/logicalAnd/... are no
  longer top-level (use operators or Tensor.add etc.)
@benikm91
benikm91 force-pushed the make-functions-for-extensions branch from 1b9b9ae to 8b75f9a Compare September 27, 2026 06:17
@benikm91
benikm91 merged commit 4812716 into dimwit-dev:main Sep 27, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants