VpfK dZddlmZddlZddlZddlmZmZddlZddl Z ddl Z ddl Z ddl Z ddl Z ddlmZmZddlZddlmZddlmZddlmZdd lmZdd lmZdd lmZdd lmZdd lmZddlmZmZmZddlm Z ddlm!Z"ddlm#Z#ddlm$Z$ddlm%Z%ddlm&Z&ddlm'Z(ddlm)Z)ddlm*Z*ddlm+Z+ddlm,Z,ddlm-Z.ddl/m0Z0ddl/m1Z1ddl2m3Z3ddl4Z5ej6dej7dd d!"Z8ej9d#ej:d$e;d%d&e;d'(Z<ej6d)ej7d*d d+"Z=ej6d,ej7d-d.d/"Z>e j?e@ZAdd4ZBdd7ZCe0jDjEZFe0jGZGe0jHZIe0jJZJe0jKZLe0jMZNeZOGd8d9ejPZQdd deQjRd:d;ZSdd dddeQjRd<d=ZTdd deQjRd>ddAZUGdBdCZVdd dd eQjRdDddEZWe jXZYdddFdGdHZZddJZ[ ej\dKZ]d.e]_^ej_`e]ddOZae]beadPZce]decdQZedRdSddVZfe#jge]efdRdSdd]Zhe jie]ehd^_d.dddRd d`daZjddeZkddfZldgZmdhZnenejoe]<diZpepejqe]<djZrerejse]<ddoZtddrZudd~ZvddZwej\dZxd.ex_^exddexbde#jgexdde&_yGddezZ{e{Z|GddZ}e}Z~ddZddZe;d%fddZdddZdZdZeeSfeeTfeeUfeefeefdZddlZejr eSZeTZeUZeZeZnddlmZee@eZ[[dS)aUPrimitives for calling Python functions on the host from JAX accelerator code. .. warning:: The host_callback APIs are deprecated as of March 20, 2024. The functionality is subsumed by the `new JAX external callbacks `_ See https://github.com/google/jax/issues/20385. This module introduces the host callback functions :func:`call`, :func:`id_tap`, and :func:`id_print`, that send their arguments from the device to the host and invoke user-defined Python functions on the host, optionally returning results back to the device computation. We show below how these functions can be used. We start with :func:`call`, and we discuss examples of calling from JAX to arbitrary Python functions on the CPU, e.g., to use NumPy CPU custom kernels. Then we show uses of :func:`id_tap` and :func:`id_print`, which have the restriction that they cannot return values from the host to the device. These primitives are generally faster because they are executed asynchronously with the device code. In particular, they can be used to tap into and to debug JAX code. Using :func:`call` to call a host function and return results to device ----------------------------------------------------------------------- Use :func:`call` to invoke a computation on the host and return NumPy arrays to the device computation. Host computation is useful, e.g., when a device computation needs some data that requires I/O on the host, or it needs a library that is available on the host and you do not want to code it in JAX. For example, eigen decomposition for general matrices in JAX does not work on TPU. We can call the Numpy implementation from any JAX accelerator computation, using a host computation:: # This function runs on the host def host_eig(m: np.ndarray) -> np.ndarray: return np.linalg.eigvals(m) # This function is used in JAX def device_fun(m): # We send "m" to the host, asking it to call "host_eig" and return the result. # We have to specify the result shape and dtype, either in the form of an # example return value or any object that has `shape` and `dtype` attributes, # e.g., a NumPy array or a `jax.ShapeDtypeStruct`. return hcb.call(host_eig, m, # Given an input of shape (..., d, d), eig output has shape (..., d) result_shape=jax.ShapeDtypeStruct(m.shape[:-1], m.dtype)) The :func:`call` function and the Python host function both take a single argument and return a single result, but those can be pytrees. Note that we must tell the :func:`call` what shape and dtype to expect from the host invocation, using the ``result_shape`` keyword argument. This is important because the device code is compiled with that expectation. There will be an error raised at runtime if the actual invocation produces a different result shape. In general, **such errors and also exceptions raised by the host computation may be difficult to debug**. See the Debugging section below. This is a problem for :func:`call` but not for :func:`id_tap` because for the latter the device code does not expect a returned value. The :func:`call` API can be used inside a jit or pmap computation or inside cond/scan/while control flow. When used inside :func:`jax.pmap`, there will be separate calls to the host from each of the participating devices:: def host_sin(x, *, device): # The ``device`` argument is passed due to ``call_with_device=True`` below. print(f"Invoking host_sin with {x.shape} on {device}") return np.sin(x) # Use pmap to run the computation on two devices jax.pmap(lambda x: hcb.call(host_sin, x, result_shape=x, # Ask that the `host_sin` function be passed `device=dev` call_with_device=True))( np.ones((2, 4), dtype=np.float32)) # prints (in arbitrary order) # Invoking host_sin with (4,) on cpu:0 # Invoking host_sin with (4,) on cpu:1 Note that :func:`call` does not support any JAX transformations, but as we show below one can make use of the existing support for `Custom differentiation in JAX `_. Using :func:`id_tap` to call a Python function on the host, with no returned values ----------------------------------------------------------------------------------- The :func:`id_tap` and :func:`id_print` are special cases of :func:`call`, when you just want the side effects of your Python callback. These functions have the advantage that once the arguments have been sent to the host, the device computation can proceed without waiting for the Python callback to return. For :func:`id_tap` you can specify your Python callback to be called, while :func:`id_print` uses a built-in callback that prints the arguments to `stdout` on the host. The Python function passed to :func:`id_tap` takes two positional arguments (the value tapped from the device computation along with a ``transforms`` tuple, described below). Optionally, the function may be passed a keyword argument ``device`` with the Device from which the value was tapped. A few examples:: def host_func(arg, transforms): ...do something with arg... # calls host_func(2x, []) on host id_tap(host_func, 2 * x) # calls host_func((2x, 3x), []) id_tap(host_func, (2 * x, 3 * x)) # The argument can be a pytree # calls host_func(2x, [], device=jax.devices()[0]) id_tap(host_func, 2 * x, tap_with_device=True) # Pass the device to the tap # calls host_func(2x, [], what='activation') id_tap(functools.partial(host_func, what='activation'), 2 * x) # calls host_func(dict(x=x, y=y), what='data') id_tap(lambda tap, transforms: host_func(tap, what='data'), dict(x=x, y=y)) The above examples can all be adapted to use :func:`id_print` instead, with the difference that :func:`id_print` prints on the host the positional argument, along with any additional kwargs and the automatic kwarg ``transforms``. Using :func:`barrier_wait` to wait until all callbacks have executed -------------------------------------------------------------------- If your Python callbacks have side-effects you may need to wait until the computation has finished to ensure that the side-effects have been observed. You can use the :func:`barrier_wait` function for that purpose:: accumulator = [] def host_log(arg, transforms): # We just record the arguments in a list accumulator.append(arg) def device_fun(x): id_tap(host_log, x) id_tap(host_log, 2. * x) jax.jit(device_fun)(1.) jax.jit(device_fun)(1.) # At this point, we have started two computations, each with two # taps, but they may not have yet executed. barrier_wait() # Now we know that all the computations started before `barrier_wait` # on all devices, have finished, and all the callbacks have finished # executing. Note that :func:`barrier_wait` will start one tiny computation with one tap on each of the `jax.local_devices()` and will wait for all these taps to be received. An alternative to using :func:`barrier_wait` is to just wait for the end of the computation, if all the callbacks are :func:`call`:: accumulator = p[] def host_log(arg): # We just record the arguments in a list accumulator.append(arg) return 0. # return something def device_fun(c): y = call(host_log, x, result_shape=jax.ShapeDtypeStruct((), np.float32)) z = call(host_log, 2. * x, result_shape=jax.ShapeDtypeStruct((), np.float32)) return y + z # return something that uses both results res1 = jax.jit(device_fun)(1.) res2 = jax.jit(device_fun)(1.) res1.block_until_ready() res2.block_until_ready() Behavior under parallelization transformations ---------------------------------------------- In presence of :func:`jax.pmap` the code will run on multiple devices and each device will tap its values independently. It may be helpful to use the ``tap_with_device`` option for :func:`id_print` or :func:`id_tap`, so that you see which device is sending which data:: jax.pmap(power3, devices=jax.local_devices()[:2])(np.array([3., 4.]) # device=cpu:0 what=x,x^2: (3., 9.) # from the first device # device=cpu:1 what=x,x^2: (4., 16.) # from the second device When using :func:`jax.pmap` with multiple devices on multiple hosts, every host will receive callbacks from all of its local devices, with an operand that corresponds to each device slice. For a :func:`call`, the callback must return to each device only the slice of the result that pertains to the corresponding device. When using the experimental :func:`pjit.pjit` the code will run on multiple devices on different shards of the input. The current implementation of host callbacks will ensure that a single device will collect and outfeed the entire operand, in a single callback. The callback function is supposed to return the entire array, which will then be sent in a single infeed to the same device that issued the outfeed. This device is then responsible for sending the required shards to the other devices:: with jax.sharding.Mesh(jax.local_devices()[:2], ["d"]): pjit.pjit(power3, in_shardings=(P("d"),), out_shardings=(P("d"),))(np.array([3., 4.])) # device=TPU:0 what=x,x^2: ( [3., 4.], # [9., 16.] ) Note that the collection of the operand on one device may result in OOM if the operand was sharded across devices. When using :func:`pjit.pjit` with multiple devices on multiple hosts, only the host for the device 0 (w.r.t. the mesh) will receive the callback, with the operand collected from all participating devices on all hosts. For a :func:`call`, the callback must return the entire array for all devices on all hosts. Behavior under JAX autodiff transformations ------------------------------------------- When used under a JAX autodiff transformation, the host callback functions operate on the primal values only. Consider the following example:: def power3(x): y = x * x # Print both 'x' and 'x^2'. Must pack as a tuple. hcb.id_print((x, y), what="x,x^2") return y * x power3(3.) # what: x,x^2 : (3., 9.) (You can see these examples tested in `host_callback_test.HostCallbackTapTest.test_tap_transforms`.) When used under :func:`jax.jvp` there will be one callback with the primal values only:: jax.jvp(power3, (3.,), (0.1,)) # what: x,x^2 : (3., 9.) Similarly for :func:`jax.grad`, we get a callback from the forward computation only:: jax.grad(power3)(3.) # what: x,x^2 : (3., 9.) If you want to invoke the callback on the tangents during a :func:`jax.jvp`, you can use a custom_jvp. For example, you can define a function that does nothing interesting except that its custom_jvp will print the tangents:: @jax.custom_jvp def print_tangents(arg): return None @print_tangents.defjvp def print_tangents_jvp(primals, tangents): arg_dot, = tangents hcb.id_print(arg_dot, what="tangents") return primals, tangents Then you use this function in the places where you want to tap the tangents:: def power3_with_tangents(x): y = x * x # Print both 'x' and 'x^2'. Must pack as a tuple. hcb.id_print((x, y), what="x,x^2") print_tangents((x, y)) return y * x jax.jvp(power3_with_tangents, (3.,), (0.1,)) # what: x,x^2 : (3., 9.) # what: tangents : (0.1, 0.6) You can do a similar thing for the cotangents during :func:`jax.grad`. This time you must be careful to use in the rest of the computation the values whose cotangents you want to tap. Hence we make the ``print_cotangents`` return its argument:: @jax.custom_vjp def print_cotangents(arg): # Must return the argument for which we want the cotangent. return arg # f_fwd: a -> (b, residual) def print_cotangents_fwd(arg): return print_cotangents(arg), None # f_bwd: (residual, CT b) -> [CT a] def print_cotangents_bwd(residual, ct_b): hcb.id_print(ct_b, what="cotangents", output_stream=testing_stream) return ct_b, print_cotangents.defvjp(print_cotangents_fwd, print_cotangents_bwd) def power3_with_cotangents(x): y = x * x # Print both 'x' and 'x^2'. Must pack as a tuple. hcb.id_print((x, y), what="x,x^2", output_stream=testing_stream) (x1, y1) = print_cotangents((x, y)) # Must use the output of print_cotangents return y1 * x1 jax.grad(power3_with_cotangents)(3.) # what: x,x^2 : (3., 9.) # what: cotangents : (9., 3.) If you use :func:`ad_checkpoint.checkpoint` to rematerialize the residuals for the backward pass, then the callbacks from the primal computation will be called twice:: jax.grad(lambda x: power3(ad_checkpoint.checkpoint(power3)(x)))(3.) # what: x,x^2 : (3., 9.) # what: x,x^2 : (27., 729.) # what: x,x^2 : (3., 9.) The callbacks are, in order from: the primal computation of the inner ``power3``, the primal computation of the outer ``power3``, and the rematerialization of the residuals for the inner ``power3``. Behavior under jax.vmap ----------------------- The host callback functions :func:`id_print` and :func:`id_tap` support the vectorization transformation :func:`jax.vmap`. For :func:`jax.vmap` the arguments to the callback are batched, and the callback function is passed an additional special ``transforms`` containing a list of transformation descriptors in the form ``("batch", {"batch_dims": ...})``, where ``...``` denotes the batched dimensions for the tapped values (one entry per argument, ` `None`` denotes an argument that was broadcast). jax.vmap(power3)(np.array([2., 3.])) # transforms: [('batch', {'batch_dims': (0, 0)})] what: x,x^2 : ([2., 3.], [4., 9.]) See documentation for :func:`id_tap`, :func:`id_print`, and :func:`call`. For more usage example, see tests/host_callback_test.py. Using :func:`call` to call a TensorFlow function, with reverse-mode autodiff support ------------------------------------------------------------------------------------ Another possible use for host computation is to invoke a library written for another framework, such as TensorFlow. In this case it becomes interesting to support JAX autodiff for host callbacks by deferring to the autodiff mechanism in TensorFlow, using the :func:`jax.custom_vjp` mechanism. This is relatively easy to do, once one understands both the JAX custom VJP and the TensorFlow autodiff mechanisms. The code for how this can be done is shown in the ``call_tf_full_ad`` function in `host_callback_to_tf_test.py `_. This example supports arbitrary higher-order differentiation as well. Note that if you just want to call TensorFlow functions from JAX, you can also use the `jax2tf.call_tf function `_. Using :func:`call` to call a JAX function on another device, with reverse-mode autodiff support ------------------------------------------------------------------------------------------------ It should not be surprising that we can use host computation to invoke a JAX computation on another device. The arguments are sent from the accelerator to the host, and then to the outside device on which the JAX host computation will run, and then the results are sent back to the original accelerator. The code for how this can be done is shown in the ``call_jax_other_device function`` in `host_callback_test.py `_. Low-level details and debugging ------------------------------- The host callback functions will be executed for each device in the order in which the send operations were performed on the device. The host callback functions for multiple devices may be interleaved. The data from the devices is received by separate threads managed by the JAX runtime (one thread per device). The runtime maintains a buffer of configurable size (see the flag ``--jax_host_callback_max_queue_byte_size``). When the buffer is full, all the receiving threads are paused which eventually pauses the computation on devices. The runtime has one additional thread for each device to invoke the Python user functions with the received data. If the processing of the callbacks is slow, it may actually lead to the runtime buffer filling up, and eventually pausing the computation on the devices when they need to send something. For more details on the outfeed receiver runtime mechanism see `runtime code `_. In order to pause the execution until all data from computations already started on devices has arrived and has been processed, use :func:`barrier_wait`. Exceptions from the user-defined callback functions are logged along with their stack traces, but the receiving threads are not stopped. Instead the last exception is recorded and the subsequent :func:`barrier_wait` will raise :exc:`CallbackException` if any exception had occurred in one of the tap functions. This exception will include the text and the stack trace of the last exception encountered. One further complication arises for callback functions that must return results to the call origin device, such as :func:`call()`. This is handled differently on CPU/GPU devices compared to TPU devices. On CPU/GPU devices, in order to avoid the device computation being stuck waiting for a result that will never arrive, in case of any error during the processing of the callback (whether raised by the user-code itself or due to a mismatch of the returned value and the expected return_shape) we send the device a "fake" result of shape ``int8[12345]``. This will make the device computation abort because the received data is different than the one that it expects. On CPU the runtime will crash with a distinctive error message: ``` Check failed: buffer->length() == buffer_length (12345 vs. ...) ``` On GPU, the failure is more user-friendly and will be surfaced to the Python program as: ``` RET_CHECK failure ... Mismatch between infeed source buffer shape s8[12345] ... ``` To debug the underlying cause for these messages, see the Debugging section. On TPU devices, there is currently no shape check for infeed, so we take the safer route of not sending this fake result in case of errors. This means that the computation will hang, and no exception will be raised (but any exceptions in the callback functions will still appear in the logs). The current implementation uses the outfeed mechanism provided by XLA. The mechanism itself is quite primitive in the sense that a receiver must know exactly the shape of each incoming packet, and how many packets are expected. This makes it hard to use for multiple kinds of data in the same computation, and it is practically impossible to use it under conditionals or in loops of non-constant iteration count. Furthermore, code that uses the outfeed mechanism directly cannot be transformed by JAX. All these limitations are addressed by the host callback functions. The tapping API introduced here makes it easy to share the outfeed mechanism for multiple purposes, while supporting all transformations. **Note that after you have used the host callback functions, you cannot use lax.outfeed directly**. You may want to :func:`stop_outfeed_receiver` if you later need to use lax.outfeed. Since the actual calls to your callback functions are made from the C++ receiver, it may be hard to debug the calls. In particular, the stack trace will not include the calling code. You can use the flag ``jax_host_callback_inline`` (or the environment variable ``JAX_HOST_CALLBACK_INLINE``) to ensure that the calls to the callbacks are inlined. This works only if the calls are outside a staging context (:func:`~jax.jit` or a control-flow primitive). The C++ `receiver `_ is started automatically on the first call to :func:`id_tap`. In order to stop it properly, upon start an ``atexit`` handler is registered to call :func:`barrier_wait` with the logging name "at_exit". There are a few environment variables that you can use to turn on logging for the C++ outfeed `receiver backend `_. * ``TF_CPP_MIN_LOG_LEVEL=0``: will turn on INFO logging, needed for all below. * ``TF_CPP_MIN_VLOG_LEVEL=3``: will make all VLOG logging up to level 3 behave like INFO logs. This may be too much, but you will see which modules are logging relevant info, and then you can select which modules to log from. * ``TF_CPP_VMODULE==3`` (the module name can be either C++ or Python, without the extension). You should also use the ``--verbosity=2`` flag so that you see the logs from Python. For example, you can try to enable logging in the ``host_callback`` module: ``TF_CPP_MIN_LOG_LEVEL=0 TF_CPP_VMODULE=host_callback=3 python tests/host_callback_test.py --verbosity=2 HostCallbackIdTapTest.test_tap_jit_simple`` If you want to enable logging in lower-level implementation modules try: ``TF_CPP_MIN_LOG_LEVEL=0 TF_CPP_VMODULE=outfeed_receiver=3,host_callback=3,outfeed_receiver_py=3,outfeed_thunk=3,infeed_thunk=3,cpu_transfer_manager=3,cpu_runtime=3,xfeed_manager=3,pjrt_client=3 python tests/host_callback_test.py --verbosity=2 HostCallbackIdTapTest.test_tap_jit_simple`` (For bazel tests use --test_arg=--vmodule=... Still to do: * More performance tests. * Explore implementation with outside compilation for TPU. * Explore implementation with XLA CustomCall for CPU and GPU. ) annotationsN)CallableSequence)Anycast)api)core)config)custom_derivatives)dtypes)lax)pjit) io_callback)adbatchingpxla)mlir) partial_eval)xla) ad_checkpoint)compiler)dispatch)pretty_printer)sharding_impls)source_info_util) tree_util)util) xla_bridge) xla_client) xla_extension)hlojax_host_callback_inlineJAX_HOST_CALLBACK_INLINEFz5Inline the host_callback, if not in a staged context.)help%jax_host_callback_max_queue_byte_size%JAX_HOST_CALLBACK_MAX_QUEUE_BYTE_SIZEgAzThe size in bytes of the buffer used to hold outfeeds from each device. When this capacity is reached consuming outfeeds from the device is paused, thus potentially pausing the device computation, until the Python callback consume more outfeeds.gnA)r$ lower_boundjax_host_callback_outfeedJAX_HOST_CALLBACK_OUTFEEDzUse outfeed implementation for host_callback, even on CPU and GPU. If false, use the CustomCall implementation. Has no effect on TPU, since only the outfeed mechanism is implemented.jax_host_callback_legacyJAX_HOST_CALLBACK_LEGACYTzUse old implementation of host_callback, documented in the module docstring.If False, use the jax.experimental.io_callback implementation. See https://github.com/google/jax/issues/20385.platformstrreturnboolc"|dvp tjS)N)tpugpucudarocm)_HOST_CALLBACK_OUTFEEDvaluer,s ^/var/www/html/nettyfy-visnx/env/lib/python3.11/site-packages/jax/experimental/host_callback.py _use_outfeedr9Ds 4 4 ' &(backend xb.XlaBackendcLtj|rtddS)z;Should be called whenever outfeed (or infeed) will be used.aIhost_callback functionality isn't supported with PJRT C API. See https://jax.readthedocs.io/en/latest/debugging/index.html and https://jax.readthedocs.io/en/latest/notebooks/external_callbacks.html for alternatives. Please file a feature request at https://github.com/google/jax/issues if none of the alternatives are sufficient.N)xbusing_pjrt_c_apiNotImplementedError)r;s r8'_raise_if_using_outfeed_with_pjrt_c_apirAIs6!!    r:ceZdZdZdZdZdZdS)CallbackFlavorzSpecifies which flavor of callback to use under JAX_HOST_CALLBACK_LEGACY=False. See https://github.com/google/jax/issues/20385. N)__name__ __module__ __qualname____doc__ IO_CALLBACKPUREDEBUGr:r8rCrC^s)+ $ %%%r:rCresulttap_with_device device_indexcallback_flavorc |rd}t||0tj|\}} |D]} tj| t |||dd||} ||S| S)az Host-callback tap primitive, like identity function with a call to ``tap_func``. .. warning:: The host_callback APIs are deprecated as of March 20, 2024. The functionality is subsumed by the `new JAX external callbacks `_ See https://github.com/google/jax/issues/20385. ``id_tap`` behaves semantically like the identity function but has the side-effect that a user-defined Python function is called with the runtime value of the argument. Args: tap_func: tap function to call like ``tap_func(arg, transforms)``, with ``arg`` as described below and where ``transforms`` is the sequence of applied JAX transformations in the form ``(name, params)``. If the `tap_with_device` optional argument is True, then the invocation also includes the device from which the value is tapped as a keyword argument: ``tap_func(arg, transforms, device=dev)``. arg: the argument passed to the tap function, can be a pytree of JAX types. result: if given, specifies the return value of ``id_tap``. This value is not passed to the tap function, and in fact is not sent from the device to the host. If the ``result`` parameter is not specified then the return value of ``id_tap`` is ``arg``. tap_with_device: if True then the tap function is invoked with the device from which the tap originates as a keyword argument. device_index: specifies from which device the tap function is invoked in a SPMD program. Works only when using the outfeed implementation mechanism, i.e., does not work on CPU unless --jax_host_callback_outfeed=True. callback_flavor: if running with `JAX_HOST_CALLBACK_LEGACY=False` specifies the flavor of callback to use. See https://github.com/google/jax/issues/20385. Returns: ``arg``, or ``result`` if given. The order of execution is by data dependency: after all the arguments and the value of ``result`` if present, are computed and before the returned value is used. At least one of the returned values of ``id_tap`` must be used in the rest of the computation, or else this operation has no effect. Tapping works even for code executed on accelerators and even for code under JAX transformations. For more details see the :mod:`jax.experimental.host_callback` module documentation. zSupport for **kwargs in ``id_tap`` has been removed. Instead, pre-apply keyword arguments, either by using a closure or by passing ``functools.partial(tap_func, **kwargs)``.NT)call_with_device result_shapeidentityrRrS) TypeErrorr tree_flattenr check_arg_call) tap_funcargrPrQrRrSkwargsmsg flat_results_rcall_ress r8_deprecated_id_taprdhsn  5 C.. ,V44OL!   &%'''(  M Or:)rPrQrR output_stream thresholdrSc `tjtf||d|}t||||||S)aLike :func:`id_tap` with a printing tap function. .. warning:: The host_callback APIs are deprecated as of March 20, 2024. The functionality is subsumed by the `new JAX external callbacks `_ See https://github.com/google/jax/issues/20385. On each invocation of the printing tap, the ``kwargs`` if present will be printed first (sorted by keys). Then arg will be printed, with the arrays stringified with ``numpy.array2string``. See the :func:`id_tap` documentation. Additional keyword arguments: * ``tap_with_device`` if True, will print also the device from which the value originates. * ``output_stream`` if given then it will be used instead of the built-in ``print``. The string will be passed as ``output_stream.write(s)``. * ``threshold`` is passed to ``numpy.array2string``. * ``callback_flavor``: if running with `JAX_HOST_CALLBACK_LEGACY=False` specifies the flavor of callback to use. See https://github.com/google/jax/issues/20385. For more details see the :mod:`jax.experimental.host_callback` module documentation. )rerfrO) functoolspartial_print_tap_funcrd) r]rPrQrRrerfrSr^printers r8_deprecated_id_printrls\J  o =,9(1 = =5; = ='   %%  ' ' ''r:)rVrUrRrS callback_funcrc tjs|tjur|t dt ||||d||S)aMake a call to the host, and expect a result. .. warning:: The host_callback APIs are deprecated as of March 20, 2024. The functionality is subsumed by the `new JAX external callbacks `_ See https://github.com/google/jax/issues/20385. Args: callback_func: The Python function to invoke on the host as ``callback_func(arg)``. If the ``call_with_device`` optional argument is True, then the invocation also includes the ``device`` kwarg with the device from which the call originates: ``callback_func(arg, device=dev)``. This function must return a pytree of numpy ndarrays. arg: the argument passed to the callback function, can be a pytree of JAX types. result_shape: a value that describes the expected shape and dtype of the result. This can be a numeric scalar, from which a shape and dtype are obtained, or an object that has ``.shape`` and ``.dtype`` attributes. If the result of the callback is a pytree, then ``result_shape`` should also be a pytree with the same structure. In particular, ``result_shape`` can be `()` or `None` if the function does not have any results. The device code containing ``call`` is compiled with the expected result shape and dtype, and an error will be raised at runtime if the actual ``callback_func`` invocation returns a different kind of result. call_with_device: if True then the callback function is invoked with the device from which the call originates as a keyword argument. device_index: specifies from which device the tap function is invoked in a SPMD program. Works only when using the outfeed implementation mechanism, i.e., does not work on CPU unless --jax_host_callback_outfeed=True. callback_flavor: if running with `JAX_HOST_CALLBACK_LEGACY=False` specifies the flavor of callback to use. See https://github.com/google/jax/issues/20385. Returns: the result of the ``callback_func`` invocation. For more details see the :mod:`jax.experimental.host_callback` module documentation. NzWhen using JAX_HOST_CALLBACK_LEGACY=False you can use the `DEBUG` flavor of callback only when the `result_shape` is None. See https://github.com/google/jax/issues/20385.F)rVrUrWrRrS)_HOST_CALLBACK_LEGACYr6rCrMr@r[)rmr]rVrUrRrSs r8_deprecated_callrpsh` %---  :   }c 05(/ K K KKr:c,eZdZdZdZdZdZdZdS)_CallbackWrappercn||_||_||_tjs|rt ddSdS)NzWhen using JAX_HOST_CALLBACK_LEGACY=False, the host_callback APIs do not support `tap_with_device` and `call_with_device`. See https://github.com/google/jax/issues/20385.)rmrWrUror6r@)selfrmrWrUs r8__init__z_CallbackWrapper.__init__,sW&DDM,D &=+;=  < = ======r:cDt|j|j|jfSN)hashrmrWrUrts r8__hash__z_CallbackWrapper.__hash__6s #T]D4IJ K KKr:cb|j|jko|j|jko|j|jkSrw)rmrWrU)rtothers r8__eq__z_CallbackWrapper.__eq__9s9  %"5 5 < MU^ + <  !U%; ;=r:ctjr |j|i|S|jr||ddS|j|i|S)NrrN)ror6 _call_legacyrWrm)rtargsr^s r8__call__z_CallbackWrapper.__call__>sa"1 T  / / // /!!$q'2... T  0 0 00r:c|jr5|jr||||S|||S|jr|||S||S)Ndevice)rWrUrm)rtr]r transformss r8rz_CallbackWrapper._call_legacyGs } ' 3!!#z&!AAA!!#z222 '!!#f!555!!#&&&r:N)rGrHrIrurzr}rrrNr:r8rrrr+sb===LLL=== 111 ' ' ' ' 'r:rr)rVrUrRrWrSctjrttjt j|t j|\}}|D]} tj | i} t|||| d<|| d<|| d<|| d<|s^t j|\} } d| D} n"#t$rd|}t|wxYw| | d<t| | d <tjr=tj|i| }|s| |n||St#j|}t"j|}t|||}|t*jur&|sJt"j|||S|t*jurt#j|||| }nt7||||d }|s|n|S) N)max_callback_queue_size_bytescallbackrW arg_treedefrRc g|]=}tjtj|t j|d>S)T) canonicalize)r ShapedArraynpshaper dtype).0rbs r8 z_call..tsM888  +BHQKKaVZ9[9[9[\\888r:zresult_shape should be a pytree of values with structure matching the expected result of the callback function. The values must be either numeric scalars, or must have 'shape' and 'dtype' attributes. Got result_treedefflat_results_aval)shardingT)rordered)ror6_initialize_outfeed_receiver"_HOST_CALLBACK_MAX_QUEUE_BYTE_SIZErcheck_callablerrYrrZrr Exception ValueErrortupleoutside_call_pbind unflattenjax local_devicesrSingleDeviceShardingrCrMdebugrrL pure_callbackr)rmr]rVrUrRrWrS flat_argsrarg_paramsflat_results_shaperrr_r`callback_devicerrcs r8r[r[Ws P &H&NPPPP]###$1#66)[d t&' x(8::&&%&'& ;)2)? )M)M&88$6888 7)577c sOO  .F "'(9":":F  -!&  # #L 1 1 1{G\G\]iGjGjj')),7O|00AAH$]H%577M.... ooo i,,, j N/ / /"=,,4666hh]L#&.%)+++h$ ,88,s * B77Ci)rrerfc nd fd }|rd|D|d<|||d<dd t|D}dfd t5|r |||t |ddddS#1swxYwYdS)aZThe consumer for id_print. We provide this as a simple tapping function for printing. This is **experimental** and may not want to add many features to it; it should be easy for the user to roll their own printing function. Args: device: the device from which the value originates (only if ``tap_with_device`` was used for :func:`id_print`). output_stream: a function whose `write` method is called with the strings to be output. threshold: the value of numpy.array2string threshold parameter. **kwargs: all other keyword args are printed before printing `arg`. sr-c^|dzdSt|dS)N )writeprint)rres r8emit_strz!_print_tap_func..emit_strs6 !d(##### Ahhhhhr:c$g|] \}}|r||fn|SrNrN)rnamers r8rz#_print_tap_func..s9<<< ,f/5>T6NN$<<.s3akkakkr:r.pp.Docct|trtjtjtjdtjdtjtjfd|DtjdgSt|trtjtjtjdtjdtjtjfd|DtjdgSt|trtjtjtjdtjdtjtjfd t| Dtjd gSt|tjr(tjtj| Stjt!|S) Nz( rEc&g|] }|SrNrNrepp_vals r8rz3_print_tap_func..pp_val..!%=%=%=AffQii%=%=%=r:z )z[ c&g|] }|SrNrNrs r8rz3_print_tap_func..pp_val..rr:z ]z{ c\g|](\}}tj|d|z)S)=)pptext)rrrrs r8rz3_print_tap_func..pp_val..sG& & & +/1a"'Q'''  VVAYY && & & r:z })rf) isinstancerrgroupconcatrnestjoinbrklistdictsorteditemsrndarray array2stringr-)r]rrfs r8rz_print_tap_func..pp_vals #u Xbi   272688%=%=%=%=%=%=%=>>??  !     C   Xbi   272688%=%=%=%=%=%=%=>>??  !     C    Xbi   272688& & & & 39#))++3F3F& & &       !      C $ $ WR_SI>>> ? ?? WSXX  r:)rr-)r.r)rrr_print_tap_lockr-) r]rrrerfr^rkv_pairsrs `` @r8rjrjsf" <<<0:<<z#_values_to_avals..s7DD!t#DM!$4$455DDDDDDr:r)valss r8_values_to_avalsrs DDtDDD D DDr: outside_callargs_ape.AbstractValueSequence[pe.AbstractValue]c|rd|vsJd|vsJ|S|dJ|dJ|dJ|d}d|vr*|dr"t|dksJ||ddzS|S)NrrrR has_tokenrE)len)rWrrrs r8_outside_call_abstract_evalrs  6 ) ) ) ) f , , , , M   + + +  ! - - - # $ 0 0 001Fvk2 v;;!     vbcc{ ** r:cd|vsJtjr3|d}tj|}t ||fddi|}|St jtg|Ri|S)NrrR send_infeedF)_HOST_CALLBACK_INLINEr6r>devices_outside_call_run_callbackrapply_primitiver)rrrRrresultss r8_outside_call_implr#s} F " " " " E.)L Z\\, 'F(vSS5SFSSG N  #N DT D D DV D DDr:c|| ||i||S#|wxYw)z7Builds op_fn(*args, **kwargs) with sharding annotation.) set_shardingclear_sharding)buildersharding_protoop_fnrr^s r8_with_sharding_protor8sZ ~&&& 5$ !& ! ! Gs 3A rN)rargs_opXlaOpc |sJt|j} | s Jd|d} |d} |j| r'| s Jd|dd} t t d|} | ot| dk}| o|}d}ttj |jttj tf|||d|}tj| || |}|rt | }| }nfd |D| r|sJt$| |g}t)j}t(jjj|_d g|_|g|_t)j}t(jjj|_t9j|gt| z|gz}d | D}tj t$j|t(j |}tC||}t$"|dt$"|d }fd tGt| Dd }fd|D}n}| }||ksJd|d|d|sMt|t|ks-Jdt|dt|d||||gzS)Nz3Should be using MLIR path for `CustomCall` loweringr%The last two arguments must be tokensc"t| Srw_aval_is_emptyavals r8z0_outside_call_translation_rule..WsnT>R>R9Sr:rFrrWrc g|]I}t|ttj|j|jJSrN)rxopsConstantLiteralrzerosrr)rrcomps r8rz2_outside_call_translation_rule..lsW  $   T28DJ #C#CDDr:rDcdg|]-}tj|D]}|.SrN)raval_to_xla_shapes$with_major_to_minor_layout_if_absent)rxrs r8rz2_outside_call_translation_rule..sV-a00  4 4 6 6r:cFg|]}t|SrN)rGetTupleElement)rioutss r8rz2_outside_call_translation_rule..s9   tQ ' 'r:Tcg|];}t|rdnd.s]K ( ( G-  A   .?.C.CA.F.Fr:zgenerated_infeed (z) != send_infeed ()got  but expected . identity = )$r9r,r get_shapeis_tokenrfilterrrAr> get_backend_register_callbackrhrir_callback_handler_datareceiver add_outfeedrAfterAllr OpShardingTypeMAXIMALtypetile_assignment_dimensionstile_assignment_devices REPLICATEDrtuple_sharding_protoInfeedWithTokenShape tuple_shaperrrange) ctxavals_in avals_outrrWrRrrr use_outfeed current_tokencurrent_itokenargs_to_outfeednon_empty_flat_results_avalneed_callback_results_on_devicergenerated_infeed callback_id next_tokenr next_itokenafter_outfeed_itokenarray_sharding_prototoken_sharding_protoinfeed_sharding_protor build_infeedouts_and_tokenrr rr s @@@@r8_outside_call_translation_ruler;AsL S\**+ KKKKKK"+-2;. $  & & / / 1 1/dnn^6T6T6]6]6_6_//-///CRCL/!%V,S,S,=&?&?!@!@)1\&J%()D%E%E%I"? ?+)".*F*FGGG" $!-   +&.:: M;GG* 3#?##G KK%M #*# ,,,,!]]4.*1MNN(244","7"<"D9:56B^2'244","7"<"G!6  3'B#C#C C   !"".e &t';(<(2(8(D(DU(K(KMMl,D2G,8::n  ! !.! 4 4d((;;k899::/gg g"k [ ( ( (M+MM{MMM ) ( (  S\\S):%;%;;;;S\\->)?)?<;; J , ,,r:r(mlir.LoweringRuleContextrrWrRintc 6t|jjdkrtd|jjd}t |}|r(t jt|g|R||dS|dkrd|jjvrtd|sJ|d} |d} | j tj ks Jd | j tj ks Jd |d d} tj } | g| } tjd t jg|jd d}rg}ng}fd }t'|jjt*jt*jfr?t1j}t0jjj|_ dg|_|g|_nd }t j||| | ||d |\}}}t>j !|rtE| }| }sMt|tks-Jdt|dtdtE|||gzS)z)MLIR Lowering for `CustomCall`-based HCB.rDz)multi-platform lowering for host_callbackr)rrWrrRcpuz>The device_index feature on CPU works only when using outfeed.rrrNrNcj|^}}t|tj|fdd}rd}|S)NFrrN)rr>r)r replica_idarrays result_arraysrrWrs r8wrapped_callbackz0_outside_call_lowering..wrapped_callbacksdJ. :&+  Mm r:T)has_side_effectrrrr)#rmodule_context platformsr@r9rxla_fallback_loweringrrrr! TokenTypeget ReplicaIdOpr rruint32r)r axis_contextrSPMDAxisContextShardingContextrrrrr r!emit_python_callbackr keep_alivesappendr)r(rrWrRrrrr,r+r,r-r.rAcallback_operandscallback_operand_avalscallback_flat_results_avalrDrrr3 keep_aliver4s ` ` ` r8_outside_call_loweringrWs7   %&&** I J JJ   )! ,(X&&+L 64 %n 5 5     +!       es'9'CCC  J L LL r(-8.  s}0022 2 2 24[ 2 2 2   1 1 3 3 3 3 35\ 3 3 3"I/  *!4O4 r29%%;(+ SbS(9; 6!#!5#4!5        %%~'EF  $&&H).6HM+,#H'(4~H$$H$($=c '88X%/%/%/!':z$++J777 $?##G+  S\\S):%;%;;;;S\\->)?)?<;; g*k2 22r:r?r7)rrrrrcdd} tj||} | |} td||| || || } |rt |S|J|Jt j| \}}||krd|d|d|}t|t tj tj |}t|}td |||td tj||Ds@d|d ||d||}t||r3t t!d |}|||S#t$$r}td |||r|jdkrrtj t+jdt*jg}td|||t |ntd|||d}~wwxYw)aHPerforms the callback: callback(arg, device, transforms) Called during the device computation once we have the argument, either from an inlined callback or from an XLA computation outfeed. Returns the flat list of result arrays. If `send_infeed` then it will also send the flat list of results to the device. r.&tuple[tuple[str, dict[str, Any]], ...]c@dtfd|DS)Nc|dkr|t|dfS|dkr|tdfS|rJ|d||ifS)Nbatchr) batch_dimsmask)logical_shapes, )r)rrs r8_unpack_transformzQ_outside_call_run_callback.._unpack_transforms.._unpack_transformso TVAY///// 6>>T+++++..d..f.....Rxr:c3"K|] }|V dSrwrN)rtrbs r8rzI_outside_call_run_callback.._unpack_transforms.. s,;;1""A&;;;;;;r:r)rrbs @r8_unpack_transformsz6_outside_call_run_callback.._unpack_transformss7 ;;;; ;;; ; ;;r:z.>s_HHR##%%););)=)==HHHHHHr:z4 should have returned a result with abstract values c"t| Srwr)rbs r8rz,_outside_call_run_callback..IsnUVFWFWBWr:z#Outside call %s threw exception %s.r1i90)rzJOutside call consumer %s exception %s. Sending to infeed the error result.zDOutside call consumer %s exception %s. On TPU we do not send infeed.)r.rY)rtree_unflattenloggerrrrrYrXrsafe_maprcanonicalize_dtyperallsafe_ziprrtransfer_to_infeedrerrorr,rarangeint8)rBrrrrrWrrrrrer]unpacked_transformsresactual_flat_resultsactual_result_treedefr_canonical_flat_resultsactual_flat_results_aval non_empty_canonical_flat_resultsrs r8rrs" < < < <? [& 1 1C,,Z88 LLD+ (3 3 4 4C"$ 6]] ' ' '  * * *3<3I#3N3N00 . 0 0***,**'**nn$T]33IK^%_%_``!12H!I!I llE#V   HH#}->-E G GHHHHH[[[ **+<==[[ 5>>?WXX[[nn D+08W8W8N2P2P,Q,Q(!!"BCCC ##     LL6!DDD" E ! !"%"85PRPW9X9X9X"Y"Y!Z aq " " "!!%(>"?"?@@@@ [q " " " G' s AF!D9F I#%B9II#rrrc`|g|R}t||dd|fzS)zAdds the `transform` to the params["transforms"]. Uses a tuple representation internally, will be unpacked before the callback by _ConsumerCallable. rrN)r)rrJ)rrtransform_params new_transforms r8_add_transformrdsK +*++-  &**\266-9II L L LLr:c<tj|jdkS)Nr)mathprodrrs r8rros 4:  ! ##r:c,~tj|Srw)rinstantiate_zeros)tanr]s r8_instantiate_zerosrrs c " ""r:c d|vsJ|dstdtj|i|}t||fS)NrrWz6JVP rule is implemented only for id_tap, not for call.)r@rrr)primalstangentsrout_primals_tappeds r8_outside_call_jvp_rulervsZ F " " " "  X V W WW%*G>v>> ! " "H ,,r:cP|dstdd|vsJt|t|ksJttt||}|dd}|r |ddkrt j|it|dSJ) NrWzDdifferentiation rules are implemented only for id_tap, not for call.rrrNr)jvp transpose) r@rrmaprrJrrr)ctsrrcts_instantiatedrs r8_outside_call_transpose_rulers  f d e ee F " " " " SSYY    313==>>zz,++* /z"~11    /  - - / //,r:c |dstdd|vsJt|d|}tj|i|}||fS)NrWz=batching rules are implemented only for id_tap, not for call.rr\)r@rrr) batched_argsr]r new_paramsrvs r8_outside_call_batching_rulers`  _ ] ^ ^^ F " " " "fgz::*\8Z88# jr:cjaxprcore.ClosedJaxprhas_input_tokenhas_output_tokencbt|j||}tj||jS)z6Rewrites a ClosedJaxpr to thread the token, if needed.)_rewrite_jaxprjaxprr ClosedJaxprconsts)rrr new_jaxprs r8_rewrite_closed_jaxprrs-V\?"$5zCCC'n)oo MBRZno>>XZ ['j&'4OO) r:r core.JaxprEqnrlist[core.JaxprEqn]input_token_varcore.Varrinput_itoken_varrr(Callable[[core.AbstractValue], core.Var]c~|jtur`d|jvsJ|||j||gz|j||gzt|jddS|jtj urtj |jgd\}}} }tj |jrt|||||||dS|||j||gz|j||gzt|jt!| ddt!|dddS|jtjurtj |jd g\} |j^} } | g| ||} ||| |j||gzt|jt%d | D dS|jtjurtj |jgd \}}}}}}}}||z}|jd |||gz|j|dz} t!|dd}|jj}|d ||ddz||dz}||j|}|jj}|d ||ddz||dz}||j|}||| |jd |||gz|j|dzt|j||dz|d |dz||dzdS|jt(jurt-tj|jd}|||j||gz|j||gzt|jt1|dd|jddz|jddz|jddzdS|jt2jur}|jd}d}d|_|||j||gz|j||gzt|jt!|dd|dS|jt2jurz|jd}g|j||} d }||| |j||gzt|jt!|dd|d!d!"dS|jt:jurt-tj|jd#}|||j||gz|j||gzt|jt!|dd|jddz|jd$t@j!t@j!fz|jd%t@j!t@j!fz|jd&dz|jd'dz(dS|jtDj#urt-tj|jd#}|||j||gz|j||gzt|jt1|dddStId)|j)*aRewrite an `eqn` and append equations to `eqns`. This is only called if the current primitive uses outfeed. Assume that the current token is in `input_token_var` and the resulting token must end in `output_token_var`. Append the result of rewriting to `eqns`. rT)r)rrr cond_jaxpr cond_nconsts body_jaxpr body_nconstsNF)rrbranchesc38K|]}t|ddVdS)TN)r)rrs r8rz_rewrite_eqn..sB++*%t<<++++++r:)r) num_consts num_carryrlinearreverselengthunroll_split_transposerr)r)r)rrE)FF)rrr call_jaxprdonated_invarsin_axes)NNout_axes)rr)rrrrcJdNFzShould not be reachedrNrNr:r8unreachable_thunkz'_rewrite_eqn..unreachable_thunk9+++++r:cdSrwrNrNr:r8rz_rewrite_eqn..;sTr:)rjvp_jaxpr_thunk fun_jaxprcJdrrNrNr:r8rz'_rewrite_eqn..unreachable_thunkJrr:z illegal param)rfwd_jaxpr_thunkbwd out_treesr in_shardings out_shardings in_layouts out_layouts)rrrrrrzoutfeed rewrite )%rrrrRreplacerrrr while_pr split_dictr rr_rewrite_while_outfeed_condrcond_prscan_pr xla_pmap_prrrr custom_jvp_call_p reset_storescustom_vjp_call_jaxpr_prpjit_prr UNSPECIFIEDrremat_pr@)rrrrrrrrrarrindexoperands new_invarsrr carry_jaxprrnr_const_and_carryrnew_jaxpr_invarsnew_jaxpr_outvarsrrrrjaxpr_s r8rrs ]n$$ cj ( ( ( (KK 3:BR0S#S$'K3CEV2W$W#' d#C#C#CEEFFFFF } ###'? DDD$F$F J:q z/00 !#t_>N"24E",... fKK :2B CCK#35F"GG 0T4HH0T5IIKKK  L LMMMMM } "" ZL99IHzEHF(FOF5EFJKK s{6FHY5Z'Z ++!)+++++,,,  - -..... } ""=A_  ' ' '>(>(:J ;1a $i/A001)5++-0Z8J8K8K-LMJ%k4>>I ---.1A"##1FF+B./ 0!! (?(?GW(?(X(X!YYI!/!I+&):233)??)B,' (!! (?(?HY(?(Z(Z![[IKK [9-1ACT0UU[,- #a-a 223nDvN`NaNaGbb ddd  e e f f f f f }''dj#*\":;;JKK :2B CCK#35F"GG )*dDAA"z*:;nL 9- <J/&8:::  ; ; < < < < < }*<<< <(I,,,%1\"KK :2B CCK#35F"GG 0D$GG 1    }*BBB ;'IA3:AA0@AJ,,, KK K#35F"GG / 4FF 1$)+++  , , - - - - - } ## !3:g#6 7 7EKK :2B CCK#35F"GG +E4>>"z*:;nLJ~.%1>3MNOJ/%1>3MNO J|4|C Z 6E   * } --- $*cj1 2 2FKK :2B CCK#35F"GG $VT488    @@@ A AAr:ctj|jgd\}}} } t|dd} |j|| zd} fd| jjD} |tj |jd|| z||gz| tj t| jd| jj |j |jd}|gfd| Dz|j|jgz}tjtjg||ggt%g}t| dd}fd |jd|D}fd |j||| zD}|jd}fd | D}|j}|j}fd | D}|j}|j}|jd}|j}|j}tj ||z||gz|||gztj t|jd |j |j tj ||z||gz|||gtj t| jd| j |j g}tjd|D}tjtjg||z|gz|z||gz|g|z||gz||g} |jd}!|tj |jd|| z| dgz| z| ddz|!g|jz||gzt(jt|d| || z| j |j dS)z&Rewrite a while whose cond has outfeedrTNc0g|]}|jSrNr)rovrs r8rz/_rewrite_while_outfeed_cond..s2 jjr:r cond_before)rrc0g|]}|jSrNrrcvrs r8rz/_rewrite_while_outfeed_cond..s%JJJrzz"'22JJJr:c0g|]}|jSrNrrrrs r8rz/_rewrite_while_outfeed_cond..s2$$$jj$$$r:c0g|]}|jSrNrrs r8rz/_rewrite_while_outfeed_cond..s5$$$ j$$$r:c0g|]}|jSrNrrs r8rz/_rewrite_while_outfeed_cond..s%FFF2::bg..FFFr:c0g|]}|jSrNrrs r8rz/_rewrite_while_outfeed_cond..s%@@@RZZ((@@@r:body cond_bodyc3$K|] }|jV dSrw)r)rrs r8rz._rewrite_while_outfeed_cond..s$EE EEEEEEr:rD)rrrrrrrrRr rcall_prr source_info out_avalsrrrset join_effectsr r)"rrrrrrrrrrrtransformed_cond_jaxpr carry_invarspred1_and_token1new_cond_pred_invarnew_cond_invarsnew_cond_jaxprtransformed_body_jaxprnew_body_invars_cond_constvarsnew_body_invars_body_constvarsnew_body_invars_prednew_body_invars_carrynew_body_invars_tokennew_body_invars_itokennew_body_carry2new_body_token2new_body_itoken2new_body_pred2new_body_token3new_body_itoken3 new_body_eqnsrnew_body_jaxprpred_outs" ` r8rrs 8< jNNN8P8P4*lJ 0T4HHL<7889,$:$@$H++  *Q|^ $| 3HX6Y Y DK /5  " " " ! & . /# :#7#:;;JJJJ\JJJJz/&''z"'((**# j_':&;RGGMM.1T4HH$$$$"%*Q|^"<$$$ $$$$z,|l'BBC$$$ $J$8$;<<FFFFFFF$*_%9::%:&6&;<<@@@@<@@@/J344/Z 0 566:j21566.J344/Z 0 566  (+@ @ "8 9 : _.>? ? + /5    ! ( /    (? :oO_=` ` ?,< =t{ /5    ! ( /-(  EE}EEE F'# j4458L7MN+,/DF\.]^""_4IY7ZZ ))+- ...Z ,Q/ 0 0(++  :a |33 48H8K7L L  *122. / : #'79J&K K + ''',6  8 8 8  /       r:idc|SrwrNrs r8rrsDr:c|SrwrNr#s r8rrsTr:c|SrwrN)r(r)r*rs r8rrstr:c$t|ddS)NF)r)js r8rrsnQu&E&Er:ceZdZdZdS)CallbackExceptionzSignals that some callback function had exceptions. Raised by :func:`barrier_wait`. See the :mod:`jax.experimental.host_callback` module documentation for details. N)rGrHrIrJrNr:r8r)r)s $r:r)czeZdZUdZded<ded<ded<ded<d ed <d ed <d ed<ded<ded<dZdZdS)_CallbackHandlerDataz(Keep track of the outfeed receiver data.rrr/ initializedon_exitzthreading.Locklockztuple[Exception, str] | Nonelast_callback_exceptionztuple[XlaLocalClient, ...]clientsztuple[XlaDevice, ...]rzdict[Callable, int]consumer_registryzdict[int, Callable]consumer_registry_by_idcd|_d|_d|_tj|_d|_d|_d|_i|_ i|_ g|_ dS)NFrN) rr,r- threadingLockr.r/r0rcallback_registrycallback_registry_by_idrQrys r8ruz_CallbackHandlerData.__init__s_DMDDL  DI#'D DLDL  D#%D Dr:c>d|_d|_d|_d|_dS)z4Wait for all pending outfeeds and stop the receiver.NFrN)rr,r0rrys r8stopz_CallbackHandlerData.stops#DMDDLDLLLr:N)rGrHrIrJ__annotations__rur9rNr:r8r+r+s00------7777%%%%    ((((....$r:r+rBrcdd|D}td|||tj|}| Jd |||S#t $rG}tj}t d|||ft_ Yd}~dSd}~wwxYw)Nrac2g|]}d|j|jdS)(r)rr)ras r8rz,_callback_input_received..$s-BBBa1ag1qw111BBBr:z?Callback input received on device %s for consumer %s arrays: %sz%We should have crashed in the runtimez4Postponing exception raised in callback function: %s) rrlrrr7rJr traceback format_excrrr/)r consumer_idrB array_reprrr formatted_es r8_callback_input_receivedrD#syyBB6BBBCC*,,P K%%% # ; ? ? L L(   F   F 8FF # ## FFF&((K LLGUUU675E2222222Fs) A55 C?.Us1:::=4888888r:c6g|]}|SrNrrr;s r8rz0_initialize_outfeed_receiver..Xs$IIIg'//11IIIr:c:g|]}t|j|SrN)r9r,)rcs r8rz0_initialize_outfeed_receiver..[s'KKK!,qz2J2JKAKKKr:c6g|]}|SrNrIrJs r8rz0_initialize_outfeed_receiver..`s$VVVg'//11VVVr:rac,g|]}t|SrN)r-)rds r8rz0_initialize_outfeed_receiver..bs F F FAQ F F Fr:zBStarting outfeed_receiver for %s. max_callback_queue_size_bytes=%srDcpdt_tjsdt_t ddSdS)NTat_exit)r_on_exitrr-_deprecated_barrier_waitrNr:r8 exit_handlerz2_initialize_outfeed_receiver..exit_handlerjs<h # +,)-& +++++,,r:T)r outfeed_receiverrr.r,r>backendsrr itertoolschainr0rrArl isEnabledForloggingrMrrstartrDrrget_compile_optionsexecutable_build_optionsratexitregister) routfeed_receiver_moduler0rclients_with_outfeedclientdevices_with_outfeed device_reprrTs r8rrBsJ*:"".".)  ".".".".".".".". ::BKMM,?,?,A,A:::GIIIIIJLLG%,"%,"KKwKKK&66-f5555 G!VVAUVVVWYY   W] + +Cii F F1E F F FGG  Y"$A C C C(?(E(E "E*>$?$? '  &q! , , E)G)G% ,,, OL!!!)-&E".".".".".".".".".".".".".".".".".".sGFGG G logging_name str | Nonec tjstjdSpdtdt j t j g fdttj D]Y\}}td|tj ||}tjfd||Ztd  5fd dddn #1swxYwYtd tj.tj\}}dt_t#d ||dS) a5Blocks the calling thread until all current outfeed is processed. Waits until all callbacks from computations already running on all devices have been received and processed by the Python callbacks. Raises CallbackException if there were exceptions while processing the callbacks. This works by enqueueing a special tap computation to all devices to which we are listening for outfeed. Once all those tap computations are done, we return from barrier_wait. Note: If any of the devices are busy and cannot accept new computations, this will deadlock. Args: logging_name: an optional string that will be used in the logging statements for this invocation. See `Debugging` in the module documentation. For more details see the :mod:`jax.experimental.host_callback` module documentation. Nzbarrier_wait[%s]: start)r.ctj|}td|t j5|ttj rBfdtjD}tdt|| ddddS#1swxYwYdS)Nz9barrier_wait[%s]: at barrier_tap for device %s. Thread %scg|]}|v| SrNrN)rrOdevices_at_barriers r8rzJ_deprecated_barrier_wait..barrier_tap_received..s.???Q"#+="="= !"="="=r:z>barrier_wait[%s]: still waiting for %s devices at barrier (%s)) rrrlrr4current_threadrRrYrZrMrnotify)dev_idxrarwaiting_for_devicesrrkr.res r8barrier_tap_receivedz6_deprecated_barrier_wait..barrier_tap_receiveds6 # +G 4F LLAFI466   '''   W] + + ????*@*H??? J /002E    iikkk                  sBC!!C%(C%z1barrier_wait[%s]: enqueueing barrier on device %src$t|Srw)rd)rrps r8rz*_deprecated_barrier_wait..s()=qAAr:z'barrier_wait[%s]: waiting for callbackscXtttjkSrw)rrr)rksr8rz*_deprecated_barrier_wait..s!.//37M7U3V3VVr:zbarrier_wait[%s]: donez@There were exceptions during callback processing. Last one was: )ror6reffects_barrierrlrr4r5 Condition enumeraterrr device_putjitwait_forr/r)) red_idxrOx_on_devlast_exceptionformatted_last_exceptionrprrkr.s ` @@@@r8rSrSus(  $  F#,,,(,777   $%%%""2:;;WWheQ LLDlTUVVV~eA...HLCG A A A A!LLLXVVVV,,8,GGG XXKKVVVVWWWXXXXXXXXXXXXXXX ,,'6663?/E/],N,592  41 4 4 5 5:HI@?sD--D14D1c8tdS)aXStops the outfeed receiver runtime. .. warning:: The host_callback APIs are deprecated as of March 20, 2024. The functionality is subsumed by the `new JAX external callbacks `_ This waits for all outfeeds from computations already running on all devices, and then stops the outfeed receiver runtime. The runtime will be restarted next time you use a tap function. It should not be necessary to use this function, unless you want to start using lax.outfeed directly after having used host callbacks. N)rr9rNr:r8!_deprecated_stop_outfeed_receiverr~sr:zThe host_callback APIs are deprecated as of March 20, 2024. The functionality is subsumed by the new JAX external callbacks. See https://github.com/google/jax/issues/20385.)id_tapid_printcall barrier_waitstop_outfeed_receiver)deprecation_getattr)r,r-r.r/)r;r<)rmr)r.r)rrr.r)rr)r(r<rr/rWr/rRr=)rrrr-r.r)r.r/)rrrr/rr/r.r)rrrr/rr/r.r)rrrrrrrrrrrrrr)rrrrrrrrrrrrrr)rBr)rrr.r=)rr=rw)rerf)rJ __future__rr^enumcollections.abcrrrhrWrZrr4r?typingrrrjax._srcrr r r r r jax.experimentalrrjax._src.interpretersrrrrrperrrrrrrrrrrr> jax._src.librr jax._src.lib.mlir.dialectsr!numpyr bool_flagbool_envrint_flagint_envr=rr5ro getLoggerrGrlr9rA_xlaopsrrr%XlaShape XlaBuilderDevice XlaDeviceClientXlaLocalClientDTypeEnumrCrKrdrlrprrr[r5rrjr Primitivermultiple_resultsoutfeed_primitivesaddrdef_abstract_evalrdef_implrr;register_translationrWregister_loweringrrrrrprimitive_jvpsrprimitive_transposesrprimitive_batchersrrrrid_poutfeed_rewriterrr)TapFunctionExceptionr+rrDrrrSr~_deprecation_msg _deprecations TYPE_CHECKINGrrrrrjax._src.deprecationsr_deprecation_getattr __getattr__rNr:r8rs ffP#""""" ........  """"""!!!!!!((((((4444444444&&&&&&444444%%%%%%""""""))))))######%%%%%%%%%%%%######&&&&&&******)(FO.66 @ &5V_+FN:CC NNKK =H &&&"*)FO/77 Q )(FO.55 :   8 $ $((((        "   "      TY    )5 OOOOOh"+7.'.'.'.'.'d'3 :K:K:K:K:K:K@''''''''''''''''^ (4<-<-<-<-<-<-B!)."" $$=====@EEEE)T //"&N+++$ !<===EEE$*+++ 68a-a-a-a-a-a-H)GHHH.0 Y3Y3Y3Y3Y3Y3v~'=NNNN"TU \ \ \ \ \ ~LLLL$$$$###---%;.!&+G'/JN+4444''''TeBeBeBeBPbbbbLt~d   !!!))***KKLLLEE )########N.-// F F F F$*-Y0.0.0.0.0.f@I@I@I@I@IF   $6 !34!#78 / 0%'?@.0QR      & !( $),;OOOOOO$$X}==+ FFr: