Skip to content

Add missing, basic bool operations #154

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
benikm91 wants to merge 1 commit into
base: main
Choose a base branch
from
+125 −9
Open

Add missing, basic bool operations #154

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 28 additions & 7 deletions AGENTS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -264,8 +264,29 @@ val tanh = t.tanh

// Clipping
val clipped = t.clip(Tensor0(1.5f), Tensor0(3.5f))

// Logical operations on Bool tensors (exact shape match)
val bx = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(true, false), Array(false, true)))
val by = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(true, true), Array(false, false)))

val both = bx and by // [[true, false], [false, false]]
val either = bx or by // [[true, true], [false, true]]
val exclusive = bx xor by // [[false, true], [false, true]]
val negated = !bx // [[false, true], [true, false]]

// Broadcasting variants (! suffix, as for +! / *!)
val row = Tensor1(Axis[A]).fromArray(Array(true, false))
val bothB = bx and_! row // [[true, false], [false, false]]
val eitherB = bx or_! row // [[true, true], [false, true]]
val exclusiveB = bx xor_! row // [[false, true], [false, true]]

// Typical use: combine comparison masks
val inRange = (t > Tensor.like(t).fill(1.0f)) and (t < Tensor.like(t).fill(4.0f))
```

**Note**: `and` / `or` / `xor` are elementwise and always evaluate both operands - unlike Scala's
short-circuiting `&&` / `||` on `Boolean`, which is why those symbols are deliberately not provided.

### Reduction Operations

Reduce tensor along axis or to scalar.
Expand Down Expand Up @@ -340,7 +361,7 @@ val wrong = t.sum(Axis[C])
// dimwit.tensor.DType.Float32] in class MdocApp0 at line 53 and
// val t:
// dimwit.tensor.Tensor2[MdocApp0.this.A, MdocApp0.this.B,
// dimwit.tensor.DType.Float32] in class MdocApp0 at line 88
// dimwit.tensor.DType.Float32] in class MdocApp0 at line 99
//
```

Expand Down Expand Up @@ -368,7 +389,7 @@ val scalarBroadcast = Tensor0(5.0f).broadcastTo(tensor.shape)
val greater = tensor > Tensor0(25.0f).broadcastTo(tensor.shape)
```

**Important**: Standard operators `+`, `-`, `*`, `/` require **exact shape match**. Use `+!`, `-!`, `*!`, `/!` for broadcasting.
**Important**: Standard operators `+`, `-`, `*`, `/` require **exact shape match**. Use `+!`, `-!`, `*!`, `/!` for broadcasting. The same holds for the logical operators on `Bool` tensors: `and`, `or`, `xor` require an exact shape match, `and_!`, `or_!`, `xor_!` broadcast.

```scala
val t = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f)))
Expand All @@ -385,7 +406,7 @@ val wrong = t + 5.0f // Use +! instead
// dimwit.tensor.DType.Float32] in class MdocApp0 at line 53 and
// val t:
// dimwit.tensor.Tensor2[MdocApp0.this.A, MdocApp0.this.B,
// dimwit.tensor.DType.Float32] in class MdocApp0 at line 97
// dimwit.tensor.DType.Float32] in class MdocApp0 at line 108
//
```

Expand Down Expand Up @@ -475,19 +496,19 @@ val wrong = m1.dot(Axis[B])(m2)
// Conflicting definitions:
// val m1:
// dimwit.tensor.Tensor2[MdocApp1.this.A, MdocApp1.this.B,
// dimwit.tensor.DType.Float32] in class MdocApp1 at line 119 and
// dimwit.tensor.DType.Float32] in class MdocApp1 at line 130 and
// val m1:
// dimwit.tensor.Tensor2[MdocApp1.this.A, MdocApp1.this.B,
// dimwit.tensor.DType.Float32] in class MdocApp1 at line 122
// dimwit.tensor.DType.Float32] in class MdocApp1 at line 133
//
// error:
// Conflicting definitions:
// val m2:
// dimwit.tensor.Tensor2[MdocApp1.this.B, MdocApp1.this.C,
// dimwit.tensor.DType.Float32] in class MdocApp1 at line 120 and
// dimwit.tensor.DType.Float32] in class MdocApp1 at line 131 and
// val m2:
// dimwit.tensor.Tensor2[MdocApp1.this.C, MdocApp1.this.D,
// dimwit.tensor.DType.Float32] in class MdocApp1 at line 123
// dimwit.tensor.DType.Float32] in class MdocApp1 at line 134
//
```

Expand Down
33 changes: 32 additions & 1 deletion core/src/main/scala/dimwit/tensor/tensorops/ElementWiseOps.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -154,6 +154,19 @@ object ElementWiseOps:
// ---------------------------------------------------------
// IsBoolean operations
// ---------------------------------------------------------

/** Elementwise logical AND of two tensors of the same shape and type. */
def logicalAnd[T <: Tuple: Labels, V: IsBoolean](t1: Tensor[T, V], t2: Tensor[T, V]): Tensor[T, V] = Tensor(Jax.jnp.logical_and(t1.jaxValue, t2.jaxValue))

/** Elementwise logical OR of two tensors of the same shape and type. */
def logicalOr[T <: Tuple: Labels, V: IsBoolean](t1: Tensor[T, V], t2: Tensor[T, V]): Tensor[T, V] = Tensor(Jax.jnp.logical_or(t1.jaxValue, t2.jaxValue))

/** Elementwise logical XOR of two tensors of the same shape and type. */
def logicalXor[T <: Tuple: Labels, V: IsBoolean](t1: Tensor[T, V], t2: Tensor[T, V]): Tensor[T, V] = Tensor(Jax.jnp.logical_xor(t1.jaxValue, t2.jaxValue))

/** Elementwise logical NOT of a tensor. */
def logicalNot[T <: Tuple: Labels, V: IsBoolean](t: Tensor[T, V]): Tensor[T, V] = Tensor(Jax.jnp.logical_not(t.jaxValue))

extension [T <: Tuple: Labels, V: IsBoolean](t: Tensor[T, V])

/** returns true if all elements of the tensor are true, false otherwise */
Expand All @@ -163,4 +176,22 @@ object ElementWiseOps:
def any: Tensor0[V] = Tensor0(Jax.jnp.any(t.jaxValue))

/** returns a tensor of the same shape with each element negated (logical NOT) */
def unary_! : Tensor[T, V] = Tensor(Jax.jnp.logical_not(t.jaxValue))
def unary_! : Tensor[T, V] = logicalNot(t)

/** elementwise logical AND with another tensor of the same shape */
infix def and(other: Tensor[T, V]): Tensor[T, V] = logicalAnd(t, other)

/** elementwise logical OR with another tensor of the same shape */
infix def or(other: Tensor[T, V]): Tensor[T, V] = logicalOr(t, other)

/** elementwise logical XOR with another tensor of the same shape */
infix def xor(other: Tensor[T, V]): Tensor[T, V] = logicalXor(t, other)

/** elementwise logical AND with a broadcastable tensor */
infix def and_![O <: Tuple](other: Tensor[O, V])(using bc: Broadcast[T, O, V]): Tensor[bc.Out, V] = bc.applyTo(t, other)(logicalAnd)

/** elementwise logical OR with a broadcastable tensor */
infix def or_![O <: Tuple](other: Tensor[O, V])(using bc: Broadcast[T, O, V]): Tensor[bc.Out, V] = bc.applyTo(t, other)(logicalOr)

/** elementwise logical XOR with a broadcastable tensor */
infix def xor_![O <: Tuple](other: Tensor[O, V])(using bc: Broadcast[T, O, V]): Tensor[bc.Out, V] = bc.applyTo(t, other)(logicalXor)
43 changes: 43 additions & 0 deletions core/src/test/scala/dimwit/tensor/TensorOpsElementwiseSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,11 +83,54 @@ class TensorOpsElementwiseSuite extends DimwitTest:

describe("Boolean ops (Tensor2)"):

val c2 = Tensor2(Axis[A], Axis[B]).fromArray(
Array(
Array(true, true),
Array(false, false)
)
)

it("inverse (!)"):
(!b2) shouldEqual Tensor2(Axis[A], Axis[B]).fromArray(
Array(Array(false, true), Array(true, false))
)

it("and"):
val expected = Tensor.like(b2).fromArray(Array(true, false, false, false))
(b2 and c2) shouldEqual expected
(b2 and c2) shouldEqual (c2 and b2)

it("or"):
val expected = Tensor.like(b2).fromArray(Array(true, true, false, true))
(b2 or c2) shouldEqual expected
(b2 or c2) shouldEqual (c2 or b2)

it("xor"):
val expected = Tensor.like(b2).fromArray(Array(false, true, false, true))
(b2 xor c2) shouldEqual expected
(b2 xor c2) shouldEqual (c2 xor b2)

it("identities"):
val allTrue = Tensor.like(b2).fill(true)
val allFalse = Tensor.like(b2).fill(false)
(b2 and allTrue) shouldEqual b2
(b2 and allFalse) shouldEqual allFalse
(b2 or allFalse) shouldEqual b2
(b2 or allTrue) shouldEqual allTrue
(b2 xor allFalse) shouldEqual b2
(b2 xor allTrue) shouldEqual !b2
(b2 xor b2) shouldEqual allFalse
// De Morgan
(!(b2 and c2)) shouldEqual ((!b2) or (!c2))
(!(b2 or c2)) shouldEqual ((!b2) and (!c2))

it("broadcasting (and_! / or_! / xor_!)"):
val bA = Tensor1(Axis[A]).fromArray(Array(true, false))
(b2 and_! bA) shouldEqual Tensor.like(b2).fromArray(Array(true, false, false, false))
(b2 or_! bA) shouldEqual Tensor.like(b2).fromArray(Array(true, true, false, true))
(b2 xor_! bA) shouldEqual Tensor.like(b2).fromArray(Array(false, true, false, true))
(bA and_! b2) shouldEqual (b2 and_! bA)

describe("Casting Ops (Tensor2)"):

it("boolean casting"):
Expand Down
23 changes: 22 additions & 1 deletion mdocs/AGENTS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -204,8 +204,29 @@ val tanh = t.tanh

// Clipping
val clipped = t.clip(Tensor0(1.5f), Tensor0(3.5f))

// Logical operations on Bool tensors (exact shape match)
val bx = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(true, false), Array(false, true)))
val by = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(true, true), Array(false, false)))

val both = bx and by // [[true, false], [false, false]]
val either = bx or by // [[true, true], [false, true]]
val exclusive = bx xor by // [[false, true], [false, true]]
val negated = !bx // [[false, true], [true, false]]

// Broadcasting variants (! suffix, as for +! / *!)
val row = Tensor1(Axis[A]).fromArray(Array(true, false))
val bothB = bx and_! row // [[true, false], [false, false]]
val eitherB = bx or_! row // [[true, true], [false, true]]
val exclusiveB = bx xor_! row // [[false, true], [false, true]]

// Typical use: combine comparison masks
val inRange = (t > Tensor.like(t).fill(1.0f)) and (t < Tensor.like(t).fill(4.0f))
```

**Note**: `and` / `or` / `xor` are elementwise and always evaluate both operands - unlike Scala's
short-circuiting `&&` / `||` on `Boolean`, which is why those symbols are deliberately not provided.

### Reduction Operations

Reduce tensor along axis or to scalar.
Expand Down Expand Up @@ -276,7 +297,7 @@ val scalarBroadcast = Tensor0(5.0f).broadcastTo(tensor.shape)
val greater = tensor > Tensor0(25.0f).broadcastTo(tensor.shape)
```

**Important**: Standard operators `+`, `-`, `*`, `/` require **exact shape match**. Use `+!`, `-!`, `*!`, `/!` for broadcasting.
**Important**: Standard operators `+`, `-`, `*`, `/` require **exact shape match**. Use `+!`, `-!`, `*!`, `/!` for broadcasting. The same holds for the logical operators on `Bool` tensors: `and`, `or`, `xor` require an exact shape match, `and_!`, `or_!`, `xor_!` broadcast.

```scala mdoc:fail
val t = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f)))
Expand Down
Loading