Skip to content

Commit

Permalink
Support complex64 tensor - real scalar operations (like numpy does) (#…
Browse files Browse the repository at this point in the history 
…629)

It is always possible to operate on a Complex64 tensor with a real scalar without loss of precision, so we allow it without explicit type conversion.

This is supported by numpy, and makes complex tensor arithmetic more convenient to use (by avoiding having to explicitly make many obvious conversions to complex).

As an example, assuming `t` is a Tensor[Complex64], this let's you write expressions such as:
```
(complex(1.0 / sqrt(2.0)) *. t -. complex(1.0)) /. complex(t.size.float)
```

as follows:
```
((1.0 / sqrt(2.0)) *. t -. 1.0) /. t.size
```

which is much more easier to understand while being equally safe.
  • Loading branch information
@AngelEzquerra
AngelEzquerra authored Mar 2, 2024
1 parent 86c1342 commit d7b1fff
Show file tree
Hide file tree
Showing 2 changed files with 146 additions and 10 deletions.
133 changes: 123 additions & 10 deletions src/arraymancer/tensor/operators_broadcasted.nim
View file
Original file line number Diff line number Diff line change
Expand Up @@ -105,58 +105,58 @@ proc `/.=`*[T: SomeNumber|Complex[float32]|Complex[float64]](a: var Tensor[T], b
# # Broadcasting Tensor-Scalar and Scalar-Tensor

proc `+.`*[T: SomeNumber|Complex[float32]|Complex[float64]](val: T, t: Tensor[T]): Tensor[T] {.noinit.} =
## Broadcasted addition for tensor + scalar.
## Broadcasted addition for tensor + scalar of the same type.
returnEmptyIfEmpty(t)
result = t.map_inline(x + val)

proc `+.`*[T: SomeNumber|Complex[float32]|Complex[float64]](t: Tensor[T], val: T): Tensor[T] {.noinit.} =
## Broadcasted addition for scalar + tensor.
## Broadcasted addition for scalar + tensor of the same type.
returnEmptyIfEmpty(t)
result = t.map_inline(x + val)

proc `-.`*[T: SomeNumber|Complex[float32]|Complex[float64]](val: T, t: Tensor[T]): Tensor[T] {.noinit.} =
## Broadcasted substraction for tensor - scalar.
## Broadcasted substraction for tensor - scalar of the same type.
returnEmptyIfEmpty(t)
result = t.map_inline(val - x)

proc `-.`*[T: SomeNumber|Complex[float32]|Complex[float64]](t: Tensor[T], val: T): Tensor[T] {.noinit.} =
## Broadcasted substraction for scalar - tensor.
## Broadcasted substraction for scalar - tensor of the same type.
returnEmptyIfEmpty(t)
result = t.map_inline(x - val)

proc `*.`*[T: SomeNumber|Complex[float32]|Complex[float64]](val: T, t: Tensor[T]): Tensor[T] {.noinit.} =
## Broadcasted multiplication for tensor * scalar.
## Broadcasted multiplication for tensor * scalar of the same type.
returnEmptyIfEmpty(t)
result = t.map_inline(x * val)

proc `*.`*[T: SomeNumber|Complex[float32]|Complex[float64]](t: Tensor[T], val: T): Tensor[T] {.noinit.} =
## Broadcasted multiplication for scalar * tensor.
## Broadcasted multiplication for scalar * tensor of the same type.
returnEmptyIfEmpty(t)
result = t.map_inline(x * val)

proc `/.`*[T: SomeNumber|Complex[float32]|Complex[float64]](val: T, t: Tensor[T]): Tensor[T] {.noinit.} =
## Broadcasted division
## Broadcasted division for scalar / tensor of the same type.
returnEmptyIfEmpty(t)
when T is SomeInteger:
result = t.map_inline(val div x)
else:
result = t.map_inline(val / x)

proc `/.`*[T: SomeNumber|Complex[float32]|Complex[float64]](t: Tensor[T], val: T): Tensor[T] {.noinit.} =
## Broadcasted division
## Broadcasted division for tensor / scalar of the same type.
returnEmptyIfEmpty(t)
when T is SomeInteger:
result = t.map_inline(x div val)
else:
result = t.map_inline(x / val)

proc `^.`*[T: SomeFloat|Complex[float32]|Complex[float64]](t: Tensor[T], exponent: T): Tensor[T] {.noinit.} =
## Compute element-wise exponentiation: tensor ^ scalar.
## Compute element-wise exponentiation: tensor ^ scalar of the same type.
returnEmptyIfEmpty(t)
result = t.map_inline pow(x, exponent)

proc `^.`*[T: SomeFloat|Complex[float32]|Complex[float64]](base: T, t: Tensor[T]): Tensor[T] {.noinit.} =
## Broadcasted exponentiation: scalar ^ tensor.
## Broadcasted exponentiation: scalar ^ tensor of the same type.
returnEmptyIfEmpty(t)
result = t.map_inline pow(base, x)

Expand Down Expand Up @@ -194,6 +194,119 @@ proc `/.=`*[T: SomeNumber|Complex[float32]|Complex[float64]](t: var Tensor[T], v
t.apply_inline(x / val)


# #################################################
# # Mixed Complex64 tensor - real scalar operations
#
# It is always possible to operate on a Complex64 tensor with a real scalar
# without loss of precission, so we allow it without explicit type conversion.
# This makes complex tensor arithmetic more convenient to use.

proc `+.`*[T: SomeNumber](val: T, t: Tensor[Complex64]): Tensor[Complex64] {.noinit, inline.} =
## Broadcasted addition for real scalar + Complex64 tensor
##
## The scalar is automatically converted to Complex64 before the operation.
let complex_val = complex(float64(val))
complex_val +. t

proc `+.`*[T: SomeNumber](t: Tensor[Complex64], val: T): Tensor[Complex64] {.noinit, inline.} =
## Broadcasted addition for real scalar + Complex64 tensor
##
## The scalar is automatically converted to Complex64 before the operation.
let complex_val = complex(float64(val))
t +. complex_val

proc `-.`*[T: SomeNumber](val: T, t: Tensor[Complex64]): Tensor[Complex64] {.noinit, inline.} =
## Broadcasted subtraction for real scalar - Complex64 tensor
##
## The scalar is automatically converted to Complex64 before the operation.
let complex_val = complex(float64(val))
complex_val -. t

proc `-.`*[T: SomeNumber](t: Tensor[Complex64], val: T): Tensor[Complex64] {.noinit, inline.} =
## Broadcasted subtraction for real scalar - Complex64 tensor
##
## The scalar is automatically converted to Complex64 before the operation.
let complex_val = complex(float64(val))
t -. complex_val

proc `*.`*[T: SomeNumber](val: T, t: Tensor[Complex64]): Tensor[Complex64] {.noinit, inline.} =
## Broadcasted multiplication for real scalar * Complex64 tensor
##
## The scalar is automatically converted to Complex64 before the operation.
let complex_val = complex(float64(val))
complex_val *. t

proc `*.`*[T: SomeNumber](t: Tensor[Complex64], val: T): Tensor[Complex64] {.noinit, inline.} =
## Broadcasted multiplication for real scalar * Complex64 tensor
##
## The scalar is automatically converted to Complex64 before the operation.
let complex_val = complex(float64(val))
t *. complex_val

proc `/.`*[T: SomeNumber](val: T, t: Tensor[Complex64]): Tensor[Complex64] {.noinit, inline.} =
## Broadcasted division for real scalar / Complex64 tensor
##
## The scalar is automatically converted to Complex64 before the operation.
let complex_val = complex(float64(val))
complex_val /. t

proc `/.`*[T: SomeNumber](t: Tensor[Complex64], val: T): Tensor[Complex64] {.noinit, inline.} =
## Broadcasted division for real scalar / Complex64 tensor
##
## The scalar is automatically converted to Complex64 before the operation.
let complex_val = complex(float64(val))
t /. complex_val

proc `^.`*[T: SomeNumber](val: T, t: Tensor[Complex64]): Tensor[Complex64] {.noinit, inline.} =
## Broadcasted exponentiation for real scalar ^ Complex64 tensor
##
## The scalar is automatically converted to Complex64 before the operation.
let complex_val = complex(float64(val))
complex_val ^. t

proc `^.`*[T: SomeNumber](t: Tensor[Complex64], val: T): Tensor[Complex64] {.noinit, inline.} =
## Broadcasted exponentiation for real scalar ^ Complex64 tensor
##
## The scalar is automatically converted to Complex64 before the operation.
let complex_val = complex(float64(val))
t ^. complex_val

proc `+.=`*[T: SomeNumber](t: var Tensor[Complex64], val: T) {.inline.} =
## Complex64 tensor in-place addition with a real scalar.
##
## The scalar is automatically converted to Complex64 before the operation.
let complex_val = complex(float64(val))
t +.= complex_val

proc `-.=`*[T: SomeNumber](t: var Tensor[Complex64], val: T) {.inline.} =
## Complex64 tensor in-place subtraction of a real scalar.
##
## The scalar is automatically converted to Complex64 before the operation.
let complex_val = complex(float64(val))
t -.= complex_val

proc `*.=`*[T: SomeNumber](t: var Tensor[Complex64], val: T) {.inline.} =
## Complex64 tensor in-place multiplication with a real scalar.
##
## The scalar is automatically converted to Complex64 before the operation.
let complex_val = complex(float64(val))
t *.= complex_val

proc `/.=`*[T: SomeNumber](t: var Tensor[Complex64], val: T) {.inline.} =
## Complex64 tensor in-place division by a real scalar.
##
## The scalar is automatically converted to Complex64 before the operation.
let complex_val = complex(float64(val))
t /.= complex_val

proc `^.=`*[T: SomeNumber](t: var Tensor[Complex64], val: T) {.inline.} =
## Complex64 tensor in-place exponentiation by a real scalar.
##
## The scalar is automatically converted to Complex64 before the operation.
let complex_val = complex(float64(val))
t ^.= complex_val


# ##############################################
# Deprecated syntax

Expand Down
23 changes: 23 additions & 0 deletions tests/tensor/test_broadcasting.nim
View file
Original file line number Diff line number Diff line change
Expand Up @@ -366,6 +366,29 @@ proc main() =
[1.0/20.0],
[1.0/30.0]].toTensor.asType(Complex[float64])

test "Implicit tensor-scalar broadcasting - Tensor[Complex64] - scalar operations":
block:
let a_c = [10, 20, 30, 40].toTensor().reshape(4,1).asType(Complex[float64])

check: complex(2.0) +. a_c +. complex(3.0) == 2 +. a_c +. 3
check: complex(2.0) -. a_c -. complex(3.0) == 2 -. a_c -. 3
check: complex(2.0) *. a_c *. complex(3.0) == 2 *. a_c *. 3
check: complex(2.0) /. a_c /. complex(3.0) == 2 /. a_c /. 3
check: complex(0.5) ^. a_c ^. complex(2.0) == 0.5 ^. a_c ^. 2

block:
var a_c = complex([10.0, 20.0, 30.0, 40.0].toTensor,
[1.0, 2.0, 3.0, 4.0].toTensor)

a_c +.= 2
a_c -.= 3.0
a_c *.= 4
a_c /.= 2.0
a_c ^.= 2
let expected = complex([320.0, 1428.0, 3328.0, 6020.0].toTensor,
[72.0, 304.0, 696.0, 1248.0].toTensor)
check: a_c.mean_absolute_error(expected) < 1e-12

test "Implicit broadcasting - Sigmoid 1 ./ (1 +. exp(-x)":
block:
proc sigmoid[T: SomeFloat](t: Tensor[T]): Tensor[T]=
Expand Down

0 comments on commit d7b1fff

Please sign in to comment.