1path: "tensorflow.tpu" 2tf_module { 3 member { 4 name: "CrossShardOptimizer" 5 mtype: "<type \'type\'>" 6 } 7 member { 8 name: "PaddingSpec" 9 mtype: "<class \'enum.EnumMeta\'>" 10 } 11 member { 12 name: "XLAOptions" 13 mtype: "<type \'type\'>" 14 } 15 member { 16 name: "experimental" 17 mtype: "<type \'module\'>" 18 } 19 member_method { 20 name: "batch_parallel" 21 argspec: "args=[\'computation\', \'inputs\', \'num_shards\', \'infeed_queue\', \'device_assignment\', \'name\', \'xla_options\'], varargs=None, keywords=None, defaults=[\'None\', \'1\', \'None\', \'None\', \'None\', \'None\'], " 22 } 23 member_method { 24 name: "bfloat16_scope" 25 argspec: "args=[\'name\'], varargs=None, keywords=None, defaults=[\'None\'], " 26 } 27 member_method { 28 name: "core" 29 argspec: "args=[\'num\'], varargs=None, keywords=None, defaults=None" 30 } 31 member_method { 32 name: "cross_replica_sum" 33 argspec: "args=[\'x\', \'group_assignment\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " 34 } 35 member_method { 36 name: "initialize_system" 37 argspec: "args=[\'embedding_config\', \'job\', \'compilation_failure_closes_chips\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'True\'], " 38 } 39 member_method { 40 name: "outside_compilation" 41 argspec: "args=[\'computation\'], varargs=args, keywords=kwargs, defaults=None" 42 } 43 member_method { 44 name: "replicate" 45 argspec: "args=[\'computation\', \'inputs\', \'infeed_queue\', \'device_assignment\', \'name\', \'maximum_shapes\', \'padding_spec\', \'xla_options\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\'], " 46 } 47 member_method { 48 name: "rewrite" 49 argspec: "args=[\'computation\', \'inputs\', \'infeed_queue\', \'device_assignment\', \'name\', \'xla_options\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\'], " 50 } 51 member_method { 52 name: "shard" 53 argspec: "args=[\'computation\', \'inputs\', \'num_shards\', \'input_shard_axes\', \'outputs_from_all_shards\', \'output_shard_axes\', \'infeed_queue\', \'device_assignment\', \'name\', \'xla_options\'], varargs=None, keywords=None, defaults=[\'None\', \'1\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\', \'None\'], " 54 } 55 member_method { 56 name: "shutdown_system" 57 argspec: "args=[\'job\'], varargs=None, keywords=None, defaults=[\'None\'], " 58 } 59} 60