# 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
                       'liverange_size', 'use_def_density', 'nr_defs_and_uses',
                       'nr_broken_hints', 'nr_urgent', 'nr_rematerializable')))
  observation_spec['progress'] = tensor_spec.BoundedTensorSpec(
      dtype=tf.float32, shape=(), name='progress', minimum=0, maximum=1)

  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='index_to_evict',
      minimum=0,
      maximum=num_registers - 1)

  return time_step_spec, action_spec
  # LINT.ThenChange(mlgo/compiler_opt/rl/google/regalloc/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 in ('mask', 'nr_urgent'):

# 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/regalloc/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

                       'liverange_size', 'use_def_density', 'nr_defs_and_uses',
                       'nr_broken_hints', 'nr_urgent', 'nr_rematerializable')))
  observation_spec['progress'] = tensor_spec.BoundedTensorSpec(
      dtype=tf.float32, shape=(), name='progress', minimum=0, maximum=1)

  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='index_to_evict',
      minimum=0,
      maximum=num_registers - 1)

  return time_step_spec, action_spec
  # LINT.ThenChange(mlgo/compiler_opt/rl/google/regalloc/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 in ('mask', 'nr_urgent'):

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