| 511 | primitive_args = [] |
| 512 | |
| 513 | def callback(args): |
| 514 | nonlocal n_in, n_out, n_const, keywords, primitives |
| 515 | t = args["type"] |
| 516 | if t == "inputs": |
| 517 | n_in = len(args["inputs"]) |
| 518 | elif args["type"] == "outputs": |
| 519 | n_out = len(args["outputs"]) |
| 520 | elif args["type"] == "keyword_inputs": |
| 521 | keywords = args["keywords"] |
| 522 | elif t == "constants": |
| 523 | n_const = len(args["constants"]) |
| 524 | elif t == "primitive": |
| 525 | primitives.append(args["name"]) |
| 526 | primitive_args.append(args["arguments"]) |
| 527 | |
| 528 | mx.export_function(callback, fn, mx.array(1.0), y=mx.array(1.0)) |
| 529 | self.assertEqual(n_in, 2) |