neural_tangents.stax.FanInProd

neural_tangents.stax.FanInProd()[source]

Fan-in product.

This layer takes a number of inputs (e.g. produced by FanOut) and elementwise-multiplies the inputs to produce a single output.

Return type:

tuple[InitFn, ApplyFn, LayerKernelFn, MaskFn]

Returns:

(init_fn, apply_fn, kernel_fn).