def __repr__(self):
args = dict(f=self.__wrapped__,
reduce=repr_function(self.reduce),
axis_name=repr(self.axis_name),
static_argnums=tuple(sorted(self.static_argnums)) or None)
args = ', '.join(f'{k}={v}' for k, v in args.items())
return f'{class_name(self)}({args})'
def __repr__(self):
args = dict(f=repr_function(self.__wrapped__.f),
noise_multiplier=self.noise_multiplier,
l2_norm_clip=self.l2_norm_clip,
microbatch=self.microbatch,
batch_axis=self.batch_axis)
args = ', '.join(f'{k}={v}' for k, v in args.items())
return f'{class_name(self)}({args})'
def f(x):
if not isinstance(x, Module) and callable(x):
return repr_function(x)
x = repr(x).split('\n')
x = [x[0]] + [' ' + y for y in x[1:]]
return '\n'.join(x)
def __repr__(self):
f = repr(self.f) if isinstance(self.f, Module) else repr_function(
self.f)
return f'{class_name(self)}(f={f}, input_argnums={self.input_argnums or None})'