allowed_batch_return_types += (tf.Tensor,)
                if config.TORCH_AVAILABLE and "torch" in sys.modules:
                    import torch

                    allowed_batch_return_types += (torch.Tensor,)
                if config.JAX_AVAILABLE and "jax" in sys.modules:
                    import jax.numpy as jnp

                    allowed_batch_return_types += (jnp.ndarray,)
                all_dict_values_are_lists = all(
                    isinstance(value, allowed_batch_return_types) for value in processed_inputs.values()
                )
                if all_dict_values_are_lists is False:
                    raise TypeError(
                        f"Provided `function` which is applied to all elements of table returns a `dict` of types {[type(x) for x in processed_inputs.values()]}. When using `batched=True`, make sure provided `function` returns a `dict` of types like `{allowed_batch_return_types}`."
                    )