Dec 28, 2023 · Dec 28, 2023 · Dec 28, 2023 · Dec 28, 2023 · Dec 28, 2023 · Dec 28, 2023
diff --git a/pgx/_src/games/go.py b/pgx/_src/games/go.py
        # update state
        x = jax.lax.cond(
            (action < self.size * self.size),
            lambda:_not_pass_move(x, action, self.size),
            lambda:_pass_move(x),
            lambda:_apply_action(x, action, self.size),
            lambda:_apply_pass(x),
        )
        # increment turns
        x = x._replace(color=(x.color + 1) % 2)
        return reward_bw


 def_pass_move(state: GameState) -> GameState:
 def_apply_pass(state: GameState) -> GameState:
    return state._replace(consecutive_pass_count=state.consecutive_pass_count + 1)


 def_not_pass_move(state: GameState, action, size) -> GameState:
 def_apply_action(state: GameState, action, size) -> GameState:
    state = state._replace(consecutive_pass_count=jnp.int32(0))
    xy = action
    num_captured_stones_before = state.num_captured_stones[state.color]
Original file line number	Diff line number	Diff line change
Expand Up		@@ -56,8 +56,8 @@ def step(self, x: GameState, action: Array) -> GameState:
		# update state
		x = jax.lax.cond(
		(action < self.size * self.size),
		lambda:_not_pass_move(x, action, self.size),
		lambda:_pass_move(x),
		lambda:_apply_action(x, action, self.size),
		lambda:_apply_pass(x),
		)
		# increment turns
		x = x._replace(color=(x.color + 1) % 2)
Expand DownExpand Up		@@ -130,11 +130,11 @@ def terminal_values(self, x: GameState):
		return reward_bw


		def_pass_move(state: GameState) -> GameState:
		def_apply_pass(state: GameState) -> GameState:
		return state._replace(consecutive_pass_count=state.consecutive_pass_count + 1)


		def_not_pass_move(state: GameState, action, size) -> GameState:
		def_apply_action(state: GameState, action, size) -> GameState:
		state = state._replace(consecutive_pass_count=jnp.int32(0))
		xy = action
		num_captured_stones_before = state.num_captured_stones[state.color]
Expand Down