|---| | [ View source on GitHub](https://github.com/tensorflow/probability/blob/v0.25.0/tensorflow_probability/python/internal/distribute_lib.py#L245-L303) | Constructs a function that broadcasts inputs over named axes. tfp.experimental.distribute.make_pbroadcast_function( fn, in_axes, out_axes, out_dtype ) Given a function `fn`, `make_pbroadcast_function` returns a new one that applies `pbroadcast` to input terms according to axis names provided in `in_axes` and `out_axes`. For each output axis in each term out the output of `fn`, inputs that do not have the output axes present are pbroadcasted before that term is computed.


| ## Args || |---|---| | `fn` | a callable to be transformed to have proadcasts at its inputs. | | `in_axes` | A structure of axis names that should match the structure of the input to `fn`. If the set of input axes for an input value does not match the output axes of a particular output value, the gradient of that output value w.r.t. the input value will be psum-ed over the axes present in the output but not the input. | | `out_axes` | A structure of axis names that should match the structure of the output of `fn`. The inputs to `fn` will be pbroadcast-ed before computing output terms according to their output axes. | | `out_dtype` | A structure of dtypes that matches the output of `fn`. |



| ## Returns || |---|---| | A new function that applies pbroadcasts to the inputs of the original function. ||