Add <!, <=!, >!, >=!

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. `<!` must be written backticked (``scalar `<!` t``) or dotted
+    // (`scalar.<!(t)`): bare infix `scalar <! t` does not parse, because the lexer reads `<!` as the start
+    // of an XML literal. The other broadcasting comparisons have no such restriction.
+    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)(less)
+    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)(lessEqual)
+    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)(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)

core/src/main/scala/dimwit/tensor/tensorops/ElementWiseOps.scala

@@ -26,21 +26,55 @@ 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 `<!` b``) or dotted (`a.<!(b)`): bare infix `a <! b` does not parse, because the
+      * lexer reads `<!` as the start of an XML literal. The other broadcasting comparisons have no such restriction.
+      */
+    def `<!`[O <: Tuple](other: Tensor[O, V])(using bc: Broadcast[T, O, V]): Tensor[bc.Out, Bool] = bc.applyTo(t, other)(less)
+
+    /** 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)(lessEqual)
+
+    /** 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)(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))
 
     /** 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])

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 (`<!`)"):
+      (iAB `<!` iA) shouldEqual bools(false, false, false, false)
+      (iA `<!` iAB) shouldEqual bools(false, true, true, true)
+      iAB.<!(iA) shouldEqual bools(false, false, false, false)
+
+    it("Less or equal (<=!)"):
+      (iAB <=! iA) shouldEqual bools(true, false, false, false)
+      (iA <=! iAB) shouldEqual bools(true, true, true, true)
+
+    it("Elementwise equality (elementEquals_!)"):
+      iAB.elementEquals_!(iA) shouldEqual bools(true, false, false, false)
+      iA.elementEquals_!(iAB) shouldEqual bools(true, false, false, false)
+
+    it("Against a scalar"):
+      (iAB >! 2) shouldEqual bools(false, false, true, true)
+      (2 >! iAB) shouldEqual bools(true, false, false, false)
+      (iAB <=! 2) shouldEqual bools(true, true, false, false)
+      (iAB `<!` 2) shouldEqual bools(true, false, false, false)
+      (2 `<!` iAB) shouldEqual bools(false, false, true, true)
+
+    it("Floats too"):
+      (tAB >! 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 `<!` iAB" shouldNot compile
+
   describe("Tensor-to-Tensor Broadcasting (complex)"):
 
     val tABCD = Tensor(Shape(Axis[A] -> 2, Axis[B] -> 2, Axis[C] -> 2, Axis[D] -> 2)).fromArray(