add `__pow__`

[nstarman]

DenseArrayValue could be raised to an integer power.

Basic test:

import jax
import jax.numpy as jnp
from quax import DenseArrayValue, quaxify

x = jnp.array([1, 2, 3], dtype=jnp.float64)
DenseArrayValue(x) ** 2
> TypeError: unsupported operand type(s) for ** or pow(): 'DenseArrayValue' and 'int'

Fair enough. Looking at ArrayValue, there's no __pow__ defined. If I add that..

class ArrayValue(Value):
    ...
    __add__ = quaxify(operator.add)
    ...
    __pow__ = quaxify(operator.pow)

DenseArrayValue(x) ** 2
> Array([1., 4., 9.], dtype=float64)

[patrick-kidger]

Thankyou! Can you add __rpow__ as well?

[nstarman]

Can do! I'm trying to think when this would be useful... Perhaps something like

2 ** quax.zero.Zero(3)

[patrick-kidger]

Alright, LGTM! Thank you for adding this. :)