def output(x, **kwargs):
out = f(params, x, **kwargs)
masked_output = utils.get_masked_array(out)
return utils.nt_tree_fn()(lambda x: x.masked_value)(masked_output)
def output(x, **kwargs): out = f(params, x, **kwargs) masked_output = utils.get_masked_array(out) return masked_output.masked_value
def _f(p): out = f(p, x, **apply_fn_kwargs) # TODO(romann): normalize properly if output is masked. out = utils.get_masked_array(out) return out.masked_value
def output(x, rng): out = f(params, x, rng=rng, **apply_fn_kwargs) masked_output = utils.get_masked_array(out) return masked_output.masked_value