class_docstr: A JAX function lowered to StableHLO.
  
  Attributes:
    fun_name: the name of the exported function, for error messages.
    in_tree: a PyTreeDef describing the tuple (args, kwargs) of the lowered JAX
        function. The actual lowering does not depend on the `in_tree`, but this
        can be used to invoke the exported function using the same argument
        structure.
    in_avals: the flat tuple of input abstract values. May contain dimension
        expressions in the shapes.
    out_tree: a PyTreeDef describing the result of the lowered JAX function.
    out_avals: the flat tuple of output abstract values. May contain dimension
        expressions in the shapes, with dimension variables among those in
        `in_avals.
    in_shardings: the flattened input shardings, as long as `in_avals`.
    out_shardings: the flattened output shardings, as long as `out_avals`.
    nr_devices: the number of devices that the module has been lowered for.
    lowering_platforms: a tuple containing at least one of 'tpu', 'cpu',
        'cuda', 'rocm'. See below for the calling convention for when
        there are multiple lowering platforms.
    ordered_effects: the ordered effects present in the serialized module.
        This is present from serialization version 9. See below for the
        calling convention in presence of ordered effects.
    unordered_effects: the unordered effects present in the serialized module.
        This is present from serialization version 9.
    mlir_module_serialized: the serialized lowered VHLO module.
    mlir_module_serialization_version: a version number for the serialized module.
        See more versioning details at https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#native-serialization-versions.
    module_kept_var_idx: the sorted indices of the arguments among `in_avals` that
        must be passed to the module. The other arguments have been dropped
        because they are not used. Same length as `in_shardings`.
    uses_shape_polymorphism: whether the `mlir_module_serialized` uses shape
        polymorphism. This may be because `in_avals` contains dimension
        variables, or due to inner calls of Exported modules that have
        dimension variables or platform index arguments. Such modules need
        shape refinement before XLA compilation.
    disabled_safety_checks: a list of descriptors of safety checks that have been
        disabled at export time. See docstring for `DisabledSafetyCheck`.
    _get_vjp: an optional function that takes the current exported function and
        returns the exported VJP function.
        The VJP function takes a flat list of arguments,
        starting with the primal arguments and followed by a cotangent argument
        for each primal output. It returns a tuple with the cotangents
        corresponding to the flattened primal inputs.
  
  Calling convention for the exported module (for latest supported version):
  
  The `mlir_module` has a `main` function that takes an optional first
  platform index argument if the module supports multiple platforms
  (`len(lowering_platforms) > 1`), followed by the token arguments corresponding
  to the ordered effects, followed by the kept array
  arguments (corresponding to `module_kept_var_idx` and `in_avals`).
  The platform index is a i32 or i64 scalar encoding the index of the current
  compilation platform into the `lowering_platforms` sequence.
  
  Inner functions use a different calling convention: an optional
  platform index argument, optional dimension variable arguments
  (scalar tensors of type i32 or i64),
  followed by optional token arguments (in presence of ordered effects),
  followed by the regular array arguments.
  The dimension arguments correspond to the dimension variables appearing in
  the `args_avals`, in sorted order of their names.
  
  Consider the lowering of a function with one array argument of type "f32[w,
  2 * h]", where "w" and "h" are two dimension variables.
  Assume that we use multi-platform lowering, and we have
  one ordered effect. The `main` function will be as follows:
  
      func public main(
            platform_index: i32 {jax.global_constant="_platform_index"},
            token_in: token,
            arg: f32[?, ?]) {
         arg_w = hlo.get_dimension_size(arg, 0)
         dim1 = hlo.get_dimension_size(arg, 1)
         arg_h = hlo.floordiv(dim1, 2)
         call _check_shape_assertions(arg)  # See below
         token = new_token()
         token_out, res = call _wrapped_jax_export_main(platform_index,
                                                        arg_h,
                                                        arg_w,
                                                        token_in,
                                                        arg)
         return token_out, res
      }
  
  The actual computation is in `_wrapped_jax_export_main`, taking also
  the values of `h` and `w` dimension variables.
  
  The signature of the `_wrapped_jax_export_main` is:
  
      func private _wrapped_jax_export_main(
          platform_index: i32 {jax.global_constant="_platform_index"},
          arg_h: i32 {jax.global_constant="h"},
          arg_w: i32 {jax.global_constant="w"},
          arg_token: stablehlo.token {jax.token=True},
          arg: f32[?, ?]) -> (stablehlo.token, ...)
  
  Prior to serialization version 9 the calling convention for effects is
  different: the `main` function does not take or return a token. Instead
  the function creates dummy tokens of type `i1[0]` and passes them to the
  `_wrapped_jax_export_main`. The `_wrapped_jax_export_main`
  takes dummy tokens of type `i1[0]` and will create internally real
  tokens to pass to the inner functions. The inner functions use real
  tokens (both before and after serialization version 9)
  
  Also starting with serialization version 9, function arguments that contain
  the platform index or the dimension variable values have a
  `jax.global_constant` string attribute whose value is the name of the
  global constant, either `_platform_index` or a dimension variable name.
  The global constant name may be empty if it is not known.
  Some global constant computations use inner functions, e.g., for
  `floor_divide`. The arguments of such functions have a `jax.global_constant`
  attribute for all attributes, meaning that the result of the function is
  also a global constant.
  
  Note that `main` contains a call to `_check_shape_assertions.
  JAX tracing assumes that `arg.shape[1]` is even, and that both `w` and `h`
  have values >= 1. We must check these constraints when we invoke the
  module. We use a special custom call `@shape_assertion` that takes
  a boolean first operand, a string `error_message` attribute that may contain
  format specifiers `{0}`, `{1}`, ..., and a variadic number of integer
  scalar operands corresponding to the format specifiers.
  
       func private _check_shape_assertions(arg: f32[?, ?]) {
         # Check that w is >= 1
         arg_w = hlo.get_dimension_size(arg, 0)
         custom_call @shape_assertion(arg_w >= 1, arg_w,
            error_message="Dimension variable 'w' must have integer value >= 1. Found {0}")
         # Check that dim1 is even
         dim1 = hlo.get_dimension_size(arg, 1)
         custom_call @shape_assertion(dim1 % 2 == 0, dim1,
            error_message="Dimension variable 'h' must have integer value >= 1. Found non-zero remainder {0}")
         # Check that h >= 1
         arg_h = hlo.floordiv(dim1, 2)
         custom_call @shape_assertion(arg_h >= 1, arg_h,
            error_message=""Dimension variable 'h' must have integer value >= 1. Found {0}")
  
  If we `call_exported` with this module we perform these checks
  statically (in `call_exported_abstract_eval`).
class_docstr_tok: ["A","JAX","function","lowered","to","StableHLO",".","Attributes",":","fun_name",":","the","name","of","the","exported","function",",","for","error","messages",".","in_tree",":","a","PyTreeDef","describing","the","tuple","(","args",",","kwargs",")","of","the","lowered","JAX","function",".","The","actual","lowering","does","not","depend","on","the","`","in_tree","`",",","but","this","can","be","used","to","invoke","the","exported","function","using","the","same","argument","structure",".","in_avals",":","the","flat","tuple","of","input","abstract","values",".","May","contain","dimension","expressions","in","the","shapes",".","out_tree",":","a","PyTreeDef","describing","the","result","of","the","lowered","JAX","function",".","out_avals",":","the","flat","tuple","of","output","abstract","values",".","May","contain","dimension","expressions","in","the","shapes",",","with","dimension","variables","among","those","in","`","in_avals",".","in_shardings",":","the","flattened","input","shardings",",","as","long","as","`","in_avals","`",".","out_shardings",":","the","flattened","output","shardings",",","as","long","as","`","out_avals","`",".","nr_devices",":","the","number","of","devices","that","the","module","has","been","lowered","for",".","lowering_platforms",":","a","tuple","containing","at","least","one","of","'tpu","'",",","'cpu","'",",","'cuda","'",",","'rocm","'",".","See","below","for","the","calling","convention","for","when","there","are","multiple","lowering","platforms",".","ordered_effects",":","the","ordered","effects","present","in","the","serialized","module",".","This","is","present","from","serialization","version","9",".","See","below","for","the","calling","convention","in","presence","of","ordered","effects",".","unordered_effects",":","the","unordered","effects","present","in","the","serialized","module",".","This","is","present","from","serialization","version","9.","mlir_module_serialized",":","the","serialized","lowered","VHLO","module",".","mlir_module_serialization_version",":","a","version","number","for","the","serialized","module",".","See","more","versioning","details","at","https",":","\/\/github.com\/google\/jax\/blob\/main\/jax\/experimental\/jax2tf\/README.md","#","native-serialization-versions",".","module_kept_var_idx",":","the","sorted","indices","of","the","arguments","among","`","in_avals","`","that","must","be","passed","to","the","module",".","The","other","arguments","have","been","dropped","because","they","are","not","used",".","Same","length","as","`","in_shardings","`",".","uses_shape_polymorphism",":","whether","the","`","mlir_module_serialized","`","uses","shape","polymorphism",".","This","may","be","because","`","in_avals","`","contains","dimension","variables",",","or","due","to","inner","calls","of","Exported","modules","that","have","dimension","variables","or","platform","index","arguments",".","Such","modules","need","shape","refinement","before","XLA","compilation",".","disabled_safety_checks",":","a","list","of","descriptors","of","safety","checks","that","have","been","disabled","at","export","time",".","See","docstring","for","`","DisabledSafetyCheck","`",".","_get_vjp",":","an","optional","function","that","takes","the","current","exported","function","and","returns","the","exported","VJP","function",".","The","VJP","function","takes","a","flat","list","of","arguments",",","starting","with","the","primal","arguments","and","followed","by","a","cotangent","argument","for","each","primal","output",".","It","returns","a","tuple","with","the","cotangents","corresponding","to","the","flattened","primal","inputs",".","Calling","convention","for","the","exported","module","(","for","latest","supported","version",")",":","The","`","mlir_module","`","has","a","`","main","`","function","that","takes","an","optional","first","platform","index","argument","if","the","module","supports","multiple","platforms","(","`","len","(","lowering_platforms",")",">","1","`",")",",","followed","by","the","token","arguments","corresponding","to","the","ordered","effects",",","followed","by","the","kept","array","arguments","(","corresponding","to","`","module_kept_var_idx","`","and","`","in_avals","`",")",".","The","platform","index","is","a","i32","or","i64","scalar","encoding","the","index","of","the","current","compilation","platform","into","the","`","lowering_platforms","`","sequence",".","Inner","functions","use","a","different","calling","convention",":","an","optional","platform","index","argument",",","optional","dimension","variable","arguments","(","scalar","tensors","of","type","i32","or","i64",")",",","followed","by","optional","token","arguments","(","in","presence","of","ordered","effects",")",",","followed","by","the","regular","array","arguments",".","The","dimension","arguments","correspond","to","the","dimension","variables","appearing","in","the","`","args_avals","`",",","in","sorted","order","of","their","names",".","Consider","the","lowering","of","a","function","with","one","array","argument","of","type","``","f32","[","w",",","2","*","h","]","''",",","where","``","w","''","and","``","h","''","are","two","dimension","variables",".","Assume","that","we","use","multi-platform","lowering",",","and","we","have","one","ordered","effect",".","The","`","main","`","function","will","be","as","follows",":","func","public","main","(","platform_index",":","i32","{","jax.global_constant=","''","_platform_index","''","}",",","token_in",":","token",",","arg",":","f32","[","?",",","?","]",")","{","arg_w","=","hlo.get_dimension_size","(","arg",",","0",")","dim1","=","hlo.get_dimension_size","(","arg",",","1",")","arg_h","=","hlo.floordiv","(","dim1",",","2",")","call","_check_shape_assertions","(","arg",")","#","See","below","token","=","new_token","(",")","token_out",",","res","=","call","_wrapped_jax_export_main","(","platform_index",",","arg_h",",","arg_w",",","token_in",",","arg",")","return","token_out",",","res","}","The","actual","computation","is","in","`","_wrapped_jax_export_main","`",",","taking","also","the","values","of","`","h","`","and","`","w","`","dimension","variables",".","The","signature","of","the","`","_wrapped_jax_export_main","`","is",":","func","private","_wrapped_jax_export_main","(","platform_index",":","i32","{","jax.global_constant=","''","_platform_index","''","}",",","arg_h",":","i32","{","jax.global_constant=","''","h","''","}",",","arg_w",":","i32","{","jax.global_constant=","''","w","''","}",",","arg_token",":","stablehlo.token","{","jax.token=True","}",",","arg",":","f32","[","?",",","?","]",")","-",">","(","stablehlo.token",",","...",")","Prior","to","serialization","version","9","the","calling","convention","for","effects","is","different",":","the","`","main","`","function","does","not","take","or","return","a","token",".","Instead","the","function","creates","dummy","tokens","of","type","`","i1","[","0","]","`","and","passes","them","to","the","`","_wrapped_jax_export_main","`",".","The","`","_wrapped_jax_export_main","`","takes","dummy","tokens","of","type","`","i1","[","0","]","`","and","will","create","internally","real","tokens","to","pass","to","the","inner","functions",".","The","inner","functions","use","real","tokens","(","both","before","and","after","serialization","version","9",")","Also","starting","with","serialization","version","9",",","function","arguments","that","contain","the","platform","index","or","the","dimension","variable","values","have","a","`","jax.global_constant","`","string","attribute","whose","value","is","the","name","of","the","global","constant",",","either","`","_platform_index","`","or","a","dimension","variable","name",".","The","global","constant","name","may","be","empty","if","it","is","not","known",".","Some","global","constant","computations","use","inner","functions",",","e.g.",",","for","`","floor_divide","`",".","The","arguments","of","such","functions","have","a","`","jax.global_constant","`","attribute","for","all","attributes",",","meaning","that","the","result","of","the","function","is","also","a","global","constant",".","Note","that","`","main","`","contains","a","call","to","`","_check_shape_assertions",".","JAX","tracing","assumes","that","`","arg.shape","[","1","]","`","is","even",",","and","that","both","`","w","`","and","`","h","`","have","values",">","=","1",".","We","must","check","these","constraints","when","we","invoke","the","module",".","We","use","a","special","custom","call","`","@","shape_assertion","`","that","takes","a","boolean","first","operand",",","a","string","`","error_message","`","attribute","that","may","contain","format","specifiers","`","{","0","}","`",",","`","{","1","}","`",",","...",",","and","a","variadic","number","of","integer","scalar","operands","corresponding","to","the","format","specifiers",".","func","private","_check_shape_assertions","(","arg",":","f32","[","?",",","?","]",")","{","#","Check","that","w","is",">","=","1","arg_w","=","hlo.get_dimension_size","(","arg",",","0",")","custom_call","@","shape_assertion","(","arg_w",">","=","1",",","arg_w",",","error_message=","''","Dimension","variable","'","w","'","must","have","integer","value",">","=","1",".","Found","{","0","}","''",")","#","Check","that","dim1","is","even","dim1","=","hlo.get_dimension_size","(","arg",",","1",")","custom_call","@","shape_assertion","(","dim1","%","2","==","0",",","dim1",",","error_message=","''","Dimension","variable","'","h","'","must","have","integer","value",">","=","1",".","Found","non-zero","remainder","{","0","}","''",")","#","Check","that","h",">","=","1","arg_h","=","hlo.floordiv","(","dim1",",","2",")","custom_call","@","shape_assertion","(","arg_h",">","=","1",",","arg_h",",","error_message=","''","''","Dimension","variable","'","h","'","must","have","integer","value",">","=","1",".","Found","{","0","}","''",")","If","we","`","call_exported","`","with","this","module","we","perform","these","checks","statically","(","in","`","call_exported_abstract_eval","`",")","."]
code: def has_vjp(self) -> bool:
      return self._get_vjp is not None
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import copy
  import dataclasses
  import functools
  import itertools
  import re
  from typing import Any, Callable, Union
  import warnings
  from absl import logging
  import numpy
  import jax
  from jax import sharding
  from jax._src import ad_util
  from jax._src import config
  from jax._src import core
  from jax._src import dispatch
  from jax._src import dtypes
  from jax._src import effects
  from jax._src.interpreters import mlir
  from jax._src.interpreters import pxla
  from jax._src.lib import xla_client
  from jax._src.lib.mlir import ir
  from jax._src.lib.mlir.dialects import hlo
  from jax._src.lib.mlir.dialects import func
  from jax._src import pjit
  from jax._src import sharding_impls
  from jax._src import source_info_util
  from jax._src import tree_util
  from jax._src import util
  from jax._src import xla_bridge
  from jax.experimental.export import _shape_poly
label_desc: Use image node_id 3 for calling the Exported obj's underlying member method code with example usage: obj.has_vjp() and returns: unknown