View source on GitHub |
tf.reduce_prod(t, axis) with faster custom gradient.
tfl.kronecker_factored_lattice_lib.custom_reduce_prod(
t, axis
)
Shows comparable speed on CPU, up to 2x speed up on GPU, and 7x on TPU.
Args | |
---|---|
t
|
The tensor to reduce. |
axis
|
The dimension to reduce. |
Returns | |
---|---|
prod(t) and grad(prod(t)) |