OpRecord = collections.namedtuple(
"OpRecord",
["name", "nargs", "dtypes", "shapes", "rng", "diff_modes", "test_name"])
def op_record(name, nargs, dtypes, shapes, rng, diff_modes, test_name=None):
test_name = test_name or name
return OpRecord(name, nargs, dtypes, shapes, rng, diff_modes, test_name)
JAX_ONE_TO_ONE_OP_RECORDS = [
op_record("abs", 1, default_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
op_record("add", 2, default_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
op_record("ceil", 1, float_dtypes, all_shapes, jtu.rand_default(), []),
op_record("conj", 1, numeric_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
op_record("conjugate", 1, numeric_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
op_record("equal", 2, default_dtypes, all_shapes, jtu.rand_some_equal(), []),
op_record("exp", 1, numeric_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
op_record("floor", 1, float_dtypes, all_shapes, jtu.rand_default(), []),
op_record("greater", 2, default_dtypes, all_shapes, jtu.rand_some_equal(), []),
op_record("greater_equal", 2, default_dtypes, all_shapes, jtu.rand_some_equal(), []),
op_record("isfinite", 1, numeric_dtypes, all_shapes, jtu.rand_some_inf(), []),
op_record("less", 2, default_dtypes, all_shapes, jtu.rand_some_equal(), []),
op_record("less_equal", 2, default_dtypes, all_shapes, jtu.rand_some_equal(), []),
op_record("log", 1, numeric_dtypes, all_shapes, jtu.rand_positive(), ["rev"]),
op_record("logical_and", 2, default_dtypes, all_shapes, jtu.rand_bool(), []),
op_record("logical_not", 1, default_dtypes, all_shapes, jtu.rand_bool(), []),
op_record("logical_or", 2, default_dtypes, all_shapes, jtu.rand_bool(), []),
op_record("logical_xor", 2, default_dtypes, all_shapes, jtu.rand_bool(), []),
op_record("maximum", 2, default_dtypes, all_shapes, jtu.rand_some_inf(), []),
op_record("minimum", 2, default_dtypes, all_shapes, jtu.rand_some_inf(), []),
op_record("multiply", 2, default_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
def op_record(name, nargs, dtypes, shapes, rng, diff_modes, test_name=None):
test_name = test_name or name
return OpRecord(name, nargs, dtypes, shapes, rng, diff_modes, test_name)
JAX_ONE_TO_ONE_OP_RECORDS = [
op_record("abs", 1, default_dtypes, all_shapes, jtu.rand_default(),
["rev"]),
op_record("add", 2, default_dtypes, all_shapes, jtu.rand_default(),
["rev"]),
op_record("ceil", 1, float_dtypes, all_shapes, jtu.rand_default(), []),
op_record("conj", 1, numeric_dtypes, all_shapes, jtu.rand_default(),
["rev"]),
op_record("conjugate", 1, numeric_dtypes, all_shapes, jtu.rand_default(),
["rev"]),
op_record("equal", 2, default_dtypes, all_shapes, jtu.rand_some_equal(),
[]),
op_record("exp", 1, numeric_dtypes, all_shapes, jtu.rand_default(),
["rev"]),
op_record("floor", 1, float_dtypes, all_shapes, jtu.rand_default(), []),
op_record("greater", 2, default_dtypes, all_shapes, jtu.rand_some_equal(),
[]),
op_record("greater_equal", 2, default_dtypes, all_shapes,
jtu.rand_some_equal(), []),
op_record("less", 2, default_dtypes, all_shapes, jtu.rand_some_equal(),
[]),
op_record("less_equal", 2, default_dtypes, all_shapes,
jtu.rand_some_equal(), []),
op_record("log", 1, numeric_dtypes, all_shapes, jtu.rand_positive(),
["rev"]),
op_record("logical_and", 2, default_dtypes, all_shapes, jtu.rand_bool(),