From b2beccda0d29bef89563b2f82b88b1076226e289 Mon Sep 17 00:00:00 2001 From: Benjamin Meyer Date: Tue, 25 Aug 2026 00:30:55 +0200 Subject: [PATCH] Add !, >=! --- .../main/scala/dimwit/tensor/ValueOps.scala | 15 +++++++ .../tensor/tensorops/ElementWiseOps.scala | 44 ++++++++++++++++--- .../tensor/TensorOpsBroadcastSuite.scala | 41 +++++++++++++++++ 3 files changed, 95 insertions(+), 5 deletions(-) diff --git a/core/src/main/scala/dimwit/tensor/ValueOps.scala b/core/src/main/scala/dimwit/tensor/ValueOps.scala index 9dcbd94e..f15a5992 100644 --- a/core/src/main/scala/dimwit/tensor/ValueOps.scala +++ b/core/src/main/scala/dimwit/tensor/ValueOps.scala @@ -2,8 +2,14 @@ package dimwit.tensor import dimwit.tensor.TensorOps.IsFloating import dimwit.tensor.TensorOps.IsNumber +import dimwit.tensor.DType.Bool import dimwit.tensor.tensorops.ElementWiseOps.add import dimwit.tensor.tensorops.ElementWiseOps.divide +import dimwit.tensor.tensorops.ElementWiseOps.equal +import dimwit.tensor.tensorops.ElementWiseOps.greater +import dimwit.tensor.tensorops.ElementWiseOps.greaterEqual +import dimwit.tensor.tensorops.ElementWiseOps.less +import dimwit.tensor.tensorops.ElementWiseOps.lessEqual import dimwit.tensor.tensorops.ElementWiseOps.multiply import dimwit.tensor.tensorops.ElementWiseOps.subtract import dimwit.tensor.tensorops.TensorOpsUtil.Broadcast @@ -31,6 +37,15 @@ object ValueOps: def *[V: IsNumber](t: Tensor0[V]): Tensor0[V] = multiply(Tensor0.likeDType(t)(scalar), t) def *![T <: Tuple: Labels, V: IsNumber](t: Tensor[T, V])(using bc: Broadcast[EmptyTuple, T, V]): Tensor[bc.Out, V] = bc.applyTo(Tensor0.likeDType(t)(scalar), t)(multiply) + // Comparing a scalar against a tensor. `![T <: Tuple: Labels, V: IsNumber](t: Tensor[T, V])(using bc: Broadcast[EmptyTuple, T, V]): Tensor[bc.Out, Bool] = bc.applyTo(Tensor0.likeDType(t)(scalar), t)(greater) + def >=![T <: Tuple: Labels, V: IsNumber](t: Tensor[T, V])(using bc: Broadcast[EmptyTuple, T, V]): Tensor[bc.Out, Bool] = bc.applyTo(Tensor0.likeDType(t)(scalar), t)(greaterEqual) + def elementEquals_![T <: Tuple: Labels, V: IsNumber](t: Tensor[T, V])(using bc: Broadcast[EmptyTuple, T, V]): Tensor[bc.Out, Bool] = bc.applyTo(Tensor0.likeDType(t)(scalar), t)(equal) + extension (scalar: Float) def /[V: IsFloating](t: Tensor0[V]): Tensor0[V] = divide(Tensor0.likeDType(t)(scalar), t) diff --git a/core/src/main/scala/dimwit/tensor/tensorops/ElementWiseOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/ElementWiseOps.scala index a6161ff1..37c98f33 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/ElementWiseOps.scala +++ b/core/src/main/scala/dimwit/tensor/tensorops/ElementWiseOps.scala @@ -26,13 +26,44 @@ object ElementWiseOps: /** Elementwise minimum of two tensors. */ def minimum[T <: Tuple: Labels, V](t1: Tensor[T, V], t2: Tensor[T, V]): Tensor[T, V] = Tensor(Jax.jnp.minimum(t1.jaxValue, t2.jaxValue)) + /** Elementwise `<` of two tensors of the same shape. */ + def less[T <: Tuple: Labels, V](t1: Tensor[T, V], t2: Tensor[T, V]): Tensor[T, Bool] = Tensor(Jax.jnp.less(t1.jaxValue, t2.jaxValue)) + + /** Elementwise `<=` of two tensors of the same shape. */ + def lessEqual[T <: Tuple: Labels, V](t1: Tensor[T, V], t2: Tensor[T, V]): Tensor[T, Bool] = Tensor(Jax.jnp.less_equal(t1.jaxValue, t2.jaxValue)) + + /** Elementwise `>` of two tensors of the same shape. */ + def greater[T <: Tuple: Labels, V](t1: Tensor[T, V], t2: Tensor[T, V]): Tensor[T, Bool] = Tensor(Jax.jnp.greater(t1.jaxValue, t2.jaxValue)) + + /** Elementwise `>=` of two tensors of the same shape. */ + def greaterEqual[T <: Tuple: Labels, V](t1: Tensor[T, V], t2: Tensor[T, V]): Tensor[T, Bool] = Tensor(Jax.jnp.greater_equal(t1.jaxValue, t2.jaxValue)) + + /** Elementwise equality of two tensors of the same shape. */ + def equal[T <: Tuple: Labels, V](t1: Tensor[T, V], t2: Tensor[T, V]): Tensor[T, Bool] = Tensor(Jax.jnp.equal(t1.jaxValue, t2.jaxValue)) + // extension methods for comparisons extension [T <: Tuple: Labels, V](t: Tensor[T, V]) - def <(other: Tensor[T, V]): Tensor[T, Bool] = Tensor(Jax.jnp.less(t.jaxValue, other.jaxValue)) - def <=(other: Tensor[T, V]): Tensor[T, Bool] = Tensor(Jax.jnp.less_equal(t.jaxValue, other.jaxValue)) - def >(other: Tensor[T, V]): Tensor[T, Bool] = Tensor(Jax.jnp.greater(t.jaxValue, other.jaxValue)) - def >=(other: Tensor[T, V]): Tensor[T, Bool] = Tensor(Jax.jnp.greater_equal(t.jaxValue, other.jaxValue)) + def <(other: Tensor[T, V]): Tensor[T, Bool] = less(t, other) + def <=(other: Tensor[T, V]): Tensor[T, Bool] = lessEqual(t, other) + def >(other: Tensor[T, V]): Tensor[T, Bool] = greater(t, other) + def >=(other: Tensor[T, V]): Tensor[T, Bool] = greaterEqual(t, other) + + /** Like [[<]], but broadcasts both sides to their common shape first. + * + * Must be written backticked (``a `]], but broadcasts both sides to their common shape first. */ + def >![O <: Tuple](other: Tensor[O, V])(using bc: Broadcast[T, O, V]): Tensor[bc.Out, Bool] = bc.applyTo(t, other)(greater) + + /** Like [[>=]], but broadcasts both sides to their common shape first. */ + def >=![O <: Tuple](other: Tensor[O, V])(using bc: Broadcast[T, O, V]): Tensor[bc.Out, Bool] = bc.applyTo(t, other)(greaterEqual) /** Checks full array equality, returns true if all elements are equal */ def ===(other: Tensor[T, V]): Tensor0[Bool] = Tensor0(Jax.jnp.array_equal(t.jaxValue, other.jaxValue)) @@ -40,7 +71,10 @@ object ElementWiseOps: /** Elementwise equality, returns a tensor of bools indicating which elements are equal */ def elementEquals(other: Tensor[T, V]): Tensor[T, Bool] = require(t.shape.dimensions == other.shape.dimensions, s"Shape mismatch: ${t.shape.dimensions} vs ${other.shape.dimensions}") - Tensor(jaxValue = Jax.jnp.equal(t.jaxValue, other.jaxValue)) + equal(t, other) + + /** Like [[elementEquals]], but broadcasts both sides to their common shape first. */ + def elementEquals_![O <: Tuple](other: Tensor[O, V])(using bc: Broadcast[T, O, V]): Tensor[bc.Out, Bool] = bc.applyTo(t, other)(equal) /** Casts the elements of this tensor to a tensor of type Bool. */ def asBool: Tensor[T, Bool] = t.asType(VType[Bool]) diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsBroadcastSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsBroadcastSuite.scala index 6b5f5962..4909e27a 100644 --- a/core/src/test/scala/dimwit/tensor/TensorOpsBroadcastSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorOpsBroadcastSuite.scala @@ -110,6 +110,47 @@ class TensorOpsBroadcastSuite extends DimwitTest: ) ) + describe("Comparison Broadcasting"): + + def bools(values: Boolean*) = Tensor(iAB.shape).fromArray(values.toArray) + + it("Greater (>!)"): + (iAB >! iA) shouldEqual bools(false, true, true, true) + (iA >! iAB) shouldEqual bools(false, false, false, false) + + it("Greater or equal (>=!)"): + (iAB >=! iA) shouldEqual bools(true, true, true, true) + (iA >=! iAB) shouldEqual bools(true, false, false, false) + + it("Less (`! 2) shouldEqual bools(false, false, true, true) + (2 >! iAB) shouldEqual bools(true, false, false, false) + (iAB <=! 2) shouldEqual bools(true, true, false, false) + (iAB `! tA) shouldEqual bools(true, true, true, true) + (tA >=! tAB) shouldEqual bools(false, false, false, false) + + it("Nothing to broadcast between equal shapes"): + "iAB >! iAB" shouldNot compile + "iAB <=! iAB" shouldNot compile + "iAB ` 2, Axis[B] -> 2, Axis[C] -> 2, Axis[D] -> 2)).fromArray(