Merge pull request #3219 from google:fix-main

PiperOrigin-RevId: 550069914

.github/analytics/get_repo_metrics.py

@@ -331,7 +331,9 @@ def main(_):
  )
  issues.get()
 
- df_issues = df_issues0 = pd.DataFrame(list(_get_issues_features(issues.raw_data)))
+ df_issues = df_issues0 = pd.DataFrame(
+ list(_get_issues_features(issues.raw_data))
+ )
  df_issues['issue_response_time'] = (
  df_issues['time_labeled_or_converted'] - df_issues['created_at']
  )
@@ -350,7 +352,9 @@ def main(_):
  prs.get()
 
  df_prs = df_prs0 = pd.DataFrame(list(_get_pr_features(prs.raw_data)))
- time_response = df_prs[['time_labeled_or_assigned', 'time_review']].min(axis=1)
+ time_response = df_prs[['time_labeled_or_assigned', 'time_review']].min(
+ axis=1
+ )
  df_prs['pr_response_time'] = time_response - df_prs['ready_for_review_at']
  df_prs['pr_resolution_time'] = (
  df_prs['time_merged_or_closed'] - df_prs['ready_for_review_at']

.pre-commit-config.yaml

@@ -13,12 +13,16 @@ repos:
  hooks:
  - id: jupytext
  args: [--sync]
+- repo: https://github.com/google/pyink
+ rev: 23.5.0
+ hooks:
+ - id: pyink
 - repo: https://github.com/pre-commit/pre-commit-hooks
  rev: v4.3.0
  hooks:
  - id: check-toml
  - id: trailing-whitespace
- exclude: ^docs/.*\.md$|^dev/.*\.py$
+ exclude: ^docs/.*\.md$
 - repo: https://github.com/kynan/nbstripout
  rev: 0.6.1
  hooks:
@@ -29,7 +33,3 @@ repos:
  --extra-keys,
  "metadata.kernelspec metadata.vscode metadata.colab cell.metadata.executionInfo.user cell.metadata.executionInfo.user_tz cell.metadata.colab",
  ]
-- repo: https://github.com/google/pyink
- rev: 23.5.0
- hooks:
- - id: pyink

dev/update_requirements.py

@@ -39,6 +39,8 @@
 Alternatively, the list can also be provided from the local environment with:
 
 python dev --versions="$(pip freeze | sed s/==/-/g) flax-0.3.6"
+
+This will use the currently installed versions of all packages.
 """
 
 import pathlib

examples/imagenet/train_test.py

@@ -55,7 +55,9 @@ def test_create_model_local(self):
 
  Uses smaller inputs than `test_create_model` to due to higher compute.
  """
- model = train.create_model(model_cls=models._ResNet1Local, half_precision=False) # pylint: disable=protected-access
+ model = train.create_model(
+ model_cls=models._ResNet1Local, half_precision=False
+ ) # pylint: disable=protected-access
  params, batch_stats = train.initialized(random.PRNGKey(0), 64, model)
  variables = {'params': params, 'batch_stats': batch_stats}
  x = random.normal(random.PRNGKey(1), (1, 64, 64, 3))

examples/lm1b/temperature_sampler_test.py

@@ -37,7 +37,9 @@ def tokens_to_logits(tokens, cache):
  logits = logits.squeeze(axis=1)
  return logits, cache
 
- new_tokens = temperature_sample(tokens, cache, tokens_to_logits, key, topk=5)
+ new_tokens = temperature_sample(
+ tokens, cache, tokens_to_logits, key, topk=5
+ )
 
  np.testing.assert_array_equal(new_tokens, [[5, 6, 7, 8]])
 

examples/lm1b/train.py

@@ -285,7 +285,9 @@ def per_host_sum_pmap(in_tree):
  host_psum = jax.pmap(lambda x: jax.lax.psum(x, "i"), "i", devices=devices)
 
  def pre_pmap(xs):
- return jax.tree_util.tree_map(lambda x: jnp.broadcast_to(x, (1,) + x.shape), xs)
+ return jax.tree_util.tree_map(
+ lambda x: jnp.broadcast_to(x, (1,) + x.shape), xs
+ )
 
  def post_pmap(xs):
  return jax.tree_util.tree_map(lambda x: x[0], xs)
@@ -525,7 +527,9 @@ def encode_strings(strs, max_len):
 
  # Shard data to devices and do a training step.
  with jax.profiler.StepTraceAnnotation("train", step_num=step):
- batch = common_utils.shard(jax.tree_util.tree_map(np.asarray, next(train_iter)))
+ batch = common_utils.shard(
+ jax.tree_util.tree_map(np.asarray, next(train_iter))
+ )
  state, metrics = p_train_step(state, batch, dropout_rng=dropout_rngs)
  train_metrics.append(metrics)
 
@@ -542,7 +546,9 @@ def encode_strings(strs, max_len):
  lr = train_metrics.pop("learning_rate").mean()
  metrics_sums = jax.tree_util.tree_map(jnp.sum, train_metrics)
  denominator = metrics_sums.pop("denominator")
- summary = jax.tree_util.tree_map(lambda x: x / denominator, metrics_sums) # pylint: disable=cell-var-from-loop
+ summary = jax.tree_util.tree_map(
+ lambda x: x / denominator, metrics_sums
+ ) # pylint: disable=cell-var-from-loop
  summary["learning_rate"] = lr
  summary["perplexity"] = jnp.clip(
  jnp.exp(summary["loss"]), a_max=1.0e4

examples/nlp_seq/train.py

@@ -42,9 +42,9 @@
 
 FLAGS = flags.FLAGS
 
-flags.DEFINE_string('model_dir', default='', help=('Directory for model data.'))
+flags.DEFINE_string('model_dir', default='', help='Directory for model data.')
 
-flags.DEFINE_string('experiment', default='xpos', help=('Experiment name.'))
+flags.DEFINE_string('experiment', default='xpos', help='Experiment name.')
 
 flags.DEFINE_integer('batch_size', default=64, help='Batch size for training.')
 
@@ -58,7 +58,7 @@
  'num_train_steps', default=75000, help='Number of train steps.'
 )
 
-flags.DEFINE_float('learning_rate', default=0.05, help=('Learning rate.'))
+flags.DEFINE_float('learning_rate', default=0.05, help='Learning rate.')
 
 flags.DEFINE_float(
  'weight_decay',
@@ -74,9 +74,9 @@
  'random_seed', default=0, help='Integer for PRNG random seed.'
 )
 
-flags.DEFINE_string('train', default='', help=('Path to training data.'))
+flags.DEFINE_string('train', default='', help='Path to training data.')
 
-flags.DEFINE_string('dev', default='', help=('Path to development data.'))
+flags.DEFINE_string('dev', default='', help='Path to development data.')
 
 
 def create_learning_rate_scheduler(
@@ -354,7 +354,9 @@ def eval_step(params, batch):
  tick = time.time()
  best_dev_score = 0
  for step, batch in zip(range(num_train_steps), train_iter):
- batch = common_utils.shard(jax.tree_util.tree_map(lambda x: x._numpy(), batch)) # pylint: disable=protected-access
+ batch = common_utils.shard(
+ jax.tree_util.tree_map(lambda x: x._numpy(), batch)
+ ) # pylint: disable=protected-access
 
  state, metrics = p_train_step(state, batch, dropout_rng=dropout_rngs)
  metrics_all.append(metrics)

examples/ppo/ppo_lib.py

@@ -195,7 +195,9 @@ def get_experience(
  sim_state = sim.conn.recv()
  sim_states.append(sim_state)
  sim_states = np.concatenate(sim_states, axis=0)
- log_probs, values = agent.policy_action(state.apply_fn, state.params, sim_states)
+ log_probs, values = agent.policy_action(
+ state.apply_fn, state.params, sim_states
+ )
  log_probs, values = jax.device_get((log_probs, values))
  probs = np.exp(np.array(log_probs))
  for i, sim in enumerate(simulators):
@@ -343,7 +345,9 @@ def train(
  score = test_episodes.policy_test(1, state.apply_fn, state.params, game)
  frames = step * config.num_agents * config.actor_steps
  summary_writer.scalar('game_score', score, frames)
- logging.info('Step %s:\nframes seen %s\nscore %s\n\n', step, frames, score)
+ logging.info(
+ 'Step %s:\nframes seen %s\nscore %s\n\n', step, frames, score
+ )
 
  # Core training code.
  alpha = (

examples/seq2seq/models.py

@@ -109,7 +109,9 @@ def __call__(
  encoding format).
  """
  # Encode inputs.
- encoder = nn.RNN(nn.LSTMCell(self.hidden_size), return_carry=True, name='encoder')
+ encoder = nn.RNN(
+ nn.LSTMCell(self.hidden_size), return_carry=True, name='encoder'
+ )
  decoder = nn.RNN(
  DecoderLSTMCell(
  decoder_inputs.shape[-1], self.teacher_force, self.vocab_size

examples/sst2/models.py

@@ -197,12 +197,16 @@ def setup(self):
 
  def __call__(self, embedded_inputs, lengths):
  # Forward LSTM.
- initial_state = self.forward_lstm.initialize_carry(embedded_inputs[:, 0].shape)
+ initial_state = self.forward_lstm.initialize_carry(
+ embedded_inputs[:, 0].shape
+ )
  _, forward_outputs = self.forward_lstm(initial_state, embedded_inputs)
 
  # Backward LSTM.
  reversed_inputs = flip_sequences(embedded_inputs, lengths)
- initial_state = self.backward_lstm.initialize_carry(reversed_inputs[:, 0].shape)
+ initial_state = self.backward_lstm.initialize_carry(
+ reversed_inputs[:, 0].shape
+ )
  _, backward_outputs = self.backward_lstm(initial_state, reversed_inputs)
  backward_outputs = flip_sequences(backward_outputs, lengths)
 

examples/wmt/decode.py

@@ -155,7 +155,9 @@ def beam_init(batch_size, beam_size, max_decode_len, cache):
  finished_seqs0 = jnp.zeros((batch_size, beam_size, max_decode_len), jnp.int32)
  finished_flags0 = jnp.zeros((batch_size, beam_size), jnp.bool_)
  # add beam dimension to attention cache pytree elements
- beam_cache0 = jax.tree_util.tree_map(lambda x: add_beam_dim(x, beam_size), cache)
+ beam_cache0 = jax.tree_util.tree_map(
+ lambda x: add_beam_dim(x, beam_size), cache
+ )
  return BeamState(
  cur_index=cur_index0,
  live_logprobs=live_logprobs0,
@@ -350,13 +352,15 @@ def beam_search_loop_body_fn(state):
  [state.finished_flags, newly_finished], axis=1
  )
  # --> [batch, beams, length], [batch, beams], [batch, beams]
- top_finished_seq, top_finished_scores, top_finished_flags = (
- gather_topk_beams(
- [finished_seqs, finished_scores, finished_flags],
- finished_scores,
- batch_size,
- beam_size,
- )
+ (
+ top_finished_seq,
+ top_finished_scores,
+ top_finished_flags,
+ ) = gather_topk_beams(
+ [finished_seqs, finished_scores, finished_flags],
+ finished_scores,
+ batch_size,
+ beam_size,
  )
 
  return BeamState(

examples/wmt/train.py

@@ -355,7 +355,9 @@ def per_host_sum_pmap(in_tree):
  host_psum = jax.pmap(lambda x: jax.lax.psum(x, "i"), "i", devices=devices)
 
  def pre_pmap(xs):
- return jax.tree_util.tree_map(lambda x: jnp.broadcast_to(x, (1,) + x.shape), xs)
+ return jax.tree_util.tree_map(
+ lambda x: jnp.broadcast_to(x, (1,) + x.shape), xs
+ )
 
  def post_pmap(xs):
  return jax.tree_util.tree_map(lambda x: x[0], xs)
@@ -626,7 +628,9 @@ def decode_tokens(toks):
 
  # Shard data to devices and do a training step.
  with jax.profiler.StepTraceAnnotation("train", step_num=step):
- batch = common_utils.shard(jax.tree_util.tree_map(np.asarray, next(train_iter)))
+ batch = common_utils.shard(
+ jax.tree_util.tree_map(np.asarray, next(train_iter))
+ )
  state, metrics = p_train_step(state, batch, dropout_rng=dropout_rngs)
  train_metrics.append(metrics)
 
@@ -643,7 +647,9 @@ def decode_tokens(toks):
  lr = train_metrics.pop("learning_rate").mean()
  metrics_sums = jax.tree_util.tree_map(jnp.sum, train_metrics)
  denominator = metrics_sums.pop("denominator")
- summary = jax.tree_util.tree_map(lambda x: x / denominator, metrics_sums) # pylint: disable=cell-var-from-loop
+ summary = jax.tree_util.tree_map(
+ lambda x: x / denominator, metrics_sums
+ ) # pylint: disable=cell-var-from-loop
  summary["learning_rate"] = lr
  summary = {"train_" + k: v for k, v in summary.items()}
  writer.write_scalars(step, summary)

flax/core/axes_scan.py

@@ -156,7 +156,9 @@ def body_fn(c, xs, init_mode=False):
  'broadcasted variable has a data dependency on the scan body.'
  )
  out_flat.append(const)
- broadcast_in, constants_out = jax.tree_util.tree_unflatten(out_tree(), out_flat)
+ broadcast_in, constants_out = jax.tree_util.tree_unflatten(
+ out_tree(), out_flat
+ )
 
  c, ys = lax.scan(
  body_fn, init, xs, length=length, reverse=reverse, unroll=unroll

flax/core/meta.py

@@ -237,7 +237,9 @@ def body(mdl, c):
 
  value: Any
  names: LogicalNames = struct.field(pytree_node=False)
- mesh: Optional[jax.sharding.Mesh] = struct.field(default=None, pytree_node=False)
+ mesh: Optional[jax.sharding.Mesh] = struct.field(
+ default=None, pytree_node=False
+ )
 
  def unbox(self, apply_constraint=True) -> Any:
  """Returns the wrapped value with the partitioning applied as a sharding constraint."""

flax/errors.py

@@ -52,7 +52,9 @@ def __init__(self):
 class FlaxError(Exception):
 
  def __init__(self, message):
- error_page = 'https://flax.readthedocs.io/en/latest/api_reference/flax.errors.html'
+ error_page = (
+ 'https://flax.readthedocs.io/en/latest/api_reference/flax.errors.html'
+ )
  module_name = self.__class__.__module__
  class_name = self.__class__.__name__
  error_msg = f'{message} ({error_page}#{module_name}.{class_name})'
@@ -706,7 +708,8 @@ def __call__(self, x):
 
  def __init__(self):
  super().__init__(
- 'Can only call init, init_with_output or apply methods on an instance of the Module class, not the Module class itself'
+ 'Can only call init, init_with_output or apply methods on an instance'
+ ' of the Module class, not the Module class itself'
  )
 
 
@@ -755,7 +758,8 @@ def __call__(self, x):
 
  def __init__(self):
  super().__init__(
- 'Trying to access a property that is accessing a non-existent attribute.'
+ 'Trying to access a property that is accessing a non-existent'
+ ' attribute.'
  )
 
 
@@ -770,7 +774,8 @@ class InvalidCheckpointError(FlaxError):
 
  def __init__(self, path, step):
  super().__init__(
- f'Trying to save an outdated checkpoint at step: "{step}" and path: "{path}".'
+ f'Trying to save an outdated checkpoint at step: "{step}" and path:'
+ f' "{path}".'
  )
 
 

flax/jax_utils.py

@@ -218,7 +218,9 @@ def transpose_out(x):
 
  def body_wrapper(c, xs):
  if keepdims:
- xs = jax.tree_util.tree_map(lambda x: x.reshape((1,) * len(axis) + x.shape), xs)
+ xs = jax.tree_util.tree_map(
+ lambda x: x.reshape((1,) * len(axis) + x.shape), xs
+ )
  xs = jax.tree_util.tree_map(transpose_out, xs)
  c, ys = body_fn(c, xs)
  if keepdims:

flax/linen/attention.py

@@ -237,7 +237,9 @@ class MultiHeadDotProductAttention(Module):
  deterministic: Optional[bool] = None
  precision: PrecisionLike = None
  kernel_init: Callable[[PRNGKey, Shape, Dtype], Array] = default_kernel_init
- bias_init: Callable[[PRNGKey, Shape, Dtype], Array] = initializers.zeros_init()
+ bias_init: Callable[[PRNGKey, Shape, Dtype], Array] = (
+ initializers.zeros_init()
+ )
  use_bias: bool = True
  attention_fn: Callable[..., Array] = dot_product_attention
  decode: bool = False
@@ -316,9 +318,12 @@ def __call__(
  'cache', 'cache_index', lambda: jnp.array(0, dtype=jnp.int32)
  )
  if is_initialized:
- *batch_dims, max_length, num_heads, depth_per_head = (
- cached_key.value.shape
- )
+ (
+ *batch_dims,
+ max_length,
+ num_heads,
+ depth_per_head,
+ ) = cached_key.value.shape
  # shape check of cached keys against query input
  expected_shape = tuple(batch_dims) + (1, num_heads, depth_per_head)
  if expected_shape != query.shape: