# Path: compiler_opt/rl/feature_ops.py
# def build_quantile_map(quantile_file_dir: str):
# def discard_fn(obs: types.Float):
# def identity_fn(obs: types.Float):
# def get_normalize_fn(quantile: List[float],
#                      with_sqrt: bool,
#                      with_z_score_normalization: bool,
#                      eps: float = 1e-8,
#                      preprocessing_fn: Optional[Callable[[types.Tensor],
#                                                          types.Float]] = None):
#   def normalize(obs: types.Float):


import gin
import tensorflow as tf
from tf_agents.specs import tensor_spec
from tf_agents.trajectories import time_step
from compiler_opt.rl import feature_ops
          'callsite_cost',
          'cold_cc_penalty',
          'last_call_to_static_bonus',
          'is_multiple_blocks',
          'nested_inlines',
          'nested_inline_cost_estimate',
          'threshold',

          # inlining_default is not used as feature in training.
          'inlining_default'))
  reward_spec = tf.TensorSpec(dtype=tf.float32, shape=(), name='reward')
  time_step_spec = time_step.time_step_spec(observation_spec, reward_spec)
  action_spec = tensor_spec.BoundedTensorSpec(
      dtype=tf.int64, shape=(), name='inlining_decision', minimum=0, maximum=1)

  return time_step_spec, action_spec
  # LINT.ThenChange(mlgo/compiler_opt/rl/google/inlining/sparse_bucket_config.pbtxt)


@gin.configurable
def get_observation_processing_layer_creator(quantile_file_dir=None,
                                             with_sqrt=True,
                                             with_z_score_normalization=True,
                                             eps=1e-8):
  """Wrapper for observation_processing_layer."""
  quantile_map = feature_ops.build_quantile_map(quantile_file_dir)

  def observation_processing_layer(obs_spec):
    """Creates the layer to process observation given obs_spec."""
    if obs_spec.name == 'inlining_default':

# Path: compiler_opt/rl/feature_ops.py
# def build_quantile_map(quantile_file_dir: str):
# def discard_fn(obs: types.Float):
# def identity_fn(obs: types.Float):
# def get_normalize_fn(quantile: List[float],
#                      with_sqrt: bool,
#                      with_z_score_normalization: bool,
#                      eps: float = 1e-8,
#                      preprocessing_fn: Optional[Callable[[types.Tensor],
#                                                          types.Float]] = None):
#   def normalize(obs: types.Float):


# Path: compiler_opt/rl/inlining/config.py
import gin
import tensorflow as tf
from tf_agents.specs import tensor_spec
from tf_agents.trajectories import time_step
from compiler_opt.rl import feature_ops

          'callsite_cost',
          'cold_cc_penalty',
          'last_call_to_static_bonus',
          'is_multiple_blocks',
          'nested_inlines',
          'nested_inline_cost_estimate',
          'threshold',

          # inlining_default is not used as feature in training.
          'inlining_default'))
  reward_spec = tf.TensorSpec(dtype=tf.float32, shape=(), name='reward')
  time_step_spec = time_step.time_step_spec(observation_spec, reward_spec)
  action_spec = tensor_spec.BoundedTensorSpec(
      dtype=tf.int64, shape=(), name='inlining_decision', minimum=0, maximum=1)

  return time_step_spec, action_spec
  # LINT.ThenChange(mlgo/compiler_opt/rl/google/inlining/sparse_bucket_config.pbtxt)


@gin.configurable
def get_observation_processing_layer_creator(quantile_file_dir=None,
                                             with_sqrt=True,
                                             with_z_score_normalization=True,
                                             eps=1e-8):
  """Wrapper for observation_processing_layer."""
  quantile_map = feature_ops.build_quantile_map(quantile_file_dir)

  def observation_processing_layer(obs_spec):
    """Creates the layer to process observation given obs_spec."""
    if obs_spec.name == 'inlining_default':

      return tf.keras.layers.Lambda(feature_ops.discard_fn)