<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Copyright 2020 The TensorFlow Probability Authors.
Licensed under the Apache License, Version 2.0 (the "License");
Step1:  JAX에서 TensorFlow 확률(TFP on JAX)
<table class="tfo-notebook-buttons" align="left">
  <td>     <a target="_blank" href="https
Step2:  TFP의 최신 야간 빌드를 사용하여 TFP on JAX를 설치할 수 있습니다.
Step3:  몇 가지 유용한 Python 라이브러리를 가져옵니다.
Step4:  또한 몇 가지 기본 JAX 기능을 가져옵니다.
Step5:  TFP on JAX 가져오기
TFP on JAX를 사용하려면 jax "기판"을 가져온 후 평소의 tfp처럼 사용하면 됩니다.
Step6:  데모
Step7:  tfd.JointDistributionCoroutine를 사용하여 모델을 정의할 수 있습니다. 가중치와 바이어스 항 모두에 표준 정규 사전 분포를 넣은 후 샘플링된 레이블을 데이터에 고정하는 target_log_prob 함수를 작성합니다.
Step8:  dist에서 샘플링하여 MCMC의 초기 상태를 생성합니다. 그런 다음 무작위 키와 초기 상태를 취하는 함수를 정의하고 No-U-Turn-Sampler(NUTS)에서 500개의 샘플을 생성할 수 있습니다. jit과 같은 JAX 변환을 사용하여 XLA로 NUTS 샘플러를 컴파일할 수 있습니다.
Step9:  샘플을 사용하여 각 가중치 세트의 예측 확률을 평균화하는 방식으로 베이지안 모델 평균화(BMA)를 수행하겠습니다.
먼저 주어진 매개변수 세트에 대해 각 클래스에 대한 확률을 생성하는 함수를 작성합니다. dist.sample_distributions를 사용하여 모델의 최종 분포를 얻을 수 있습니다.
Step10:  샘플 세트에 vmap(classifier_probs)를 수행하여 각 샘플의 예측된 클래스 확률을 얻을 수 있습니다. 그런 다음 각 샘플의 평균 정확성과 베이지안 모델 평균화의 정확성을 계산합니다.
Step11:  BMA가 오류율을 거의 3분의 1로 줄이는 것으로 보입니다!
기본 사항
TFP on JAX에는 tf.Tensor와 같은 TF 객체를 수락하는 대신 JAX 아날로그를 수락하는 TF와 동일한 API가 있습니다. 예를 들어 tf.Tensor가 이전에 입력으로 사용된 곳이라면 어디에서나 API가 AX DeviceArray를 예상합니다. TFP 메서드는 tf.Tensor를 반환하는 대신에 DeviceArray를 반환합니다. TFP on JAX는 DeviceArray 목록이나 사전과 같은 JAX 객체의 중첩 구조에서도 작동합니다.
분포
대부분의 TFP 분포는 TF 대응 부분과 매우 유사한 의미 체계로 JAX에서 지원됩니다. 또한, JAX Pytrees로 등록되어 JAX 변환 함수의 입력 및 출력이 될 수도 있습니다.
기본 분포
이 분포의 log_prob 메서드는 동일하게 작동합니다.
Step12:  한 분포로부터 샘플링하려면 PRNGKey(또는 정수 목록)에서 seed 키워드 인수로 명시적으로 전달해야 합니다. 명시적으로 시드를 전달하지 못하면 오류가 발생합니다.
Step13:  분포의 형상 의미 체계는 JAX에서 동일하게 유지됩니다. 여기서 분포는 각각 event_shape와 batch_shape를 갖게 되며 많은 샘플을 이끌어낼수록 sample_shape 차원이 더 많이 추가됩니다.
예를 들어 벡터 매개변수를 가진 tfd.MultivariateNormalDiag는 벡터 이벤트 형상과 빈 배치 형상을 갖게 됩니다.
Step14:  반면에 벡터로 매개변수화된 tfd.Normal는 스칼라 이벤트 형상과 벡터 배치 형상을 갖게 됩니다.
Step15:  샘플의 log_prob를 취하는 의미 체계는 JAX에서도 동일하게 작동합니다.
Step16:  JAX DeviceArray는 NumPy 및 Matplotlib와 같은 라이브러리와 호환되므로 샘플을 직접 플로팅 함수에 공급할 수 있습니다.
Step17:  Distribution 메서드는 JAX 변환과 호환됩니다.
Step18:  TFP 분포는 JAX pytree 노드로 등록되어 있기 때문에 분포를 입력 또는 출력으로 사용하여 함수를 작성하고 이러한 함수를 jit를 사용하여 변환할 수는 있지만 아직 vmap를 적용한 함수에 대한 인수로는 지원되지 않습니다.
Step19:  변환된 분포
샘플이 Bijector를 통해 전달된 분포와 같이 변환된 분포도 기본적으로 작동합니다(bijectors도 작동합니다! 아래 참조).
Step20:  결합 분포
TFP는 구성 요소 분포를 여러 무작위 변수에 대한 단일 분포로 결합할 수 있도록 JointDistribution를 제공합니다. 현재 TFP는 JAX에서 지원되는 3가지 핵심 변화형(JointDistributionSequential, JointDistributionNamed 및 JointDistributionCoroutine)을 제공합니다. AutoBatched 변화형도 모두 지원합니다.
Step21:  기타 분포
가우시안 프로세스도 JAX 모드에서 작동합니다!
Step22:  Hidden Markov 모델도 지원합니다.
Step23:  TensorFlow 또는 XLA 비호환성에 대한 엄격한 종속성으로 인해 PixelCNN와 같은 몇몇 분포는 아직 지원하지 않습니다.
Bijectors
현재 JAX는 대부분의 TFP Bijector를 지원하고 있습니다!
Step24:  Bijector는 jit, grad 및 vmap과 같은 JAX 변환과 호환됩니다.
Step25:  RealNVP 및 FFJORD와 같은 일부 Bijector는 아직 지원하지 않습니다.
MCMC
tfp.mcmc도 JAX로 이식했으므로 JAX에서 Hamiltonian Monte Carlo(HMC) 및 No-U-Turn-Sampler(NUTS)와 같은 알고리즘도 실행할 수 있습니다.
Step26:  TFP on TF와 달리 seed 키워드 인수를 사용하여 PRNGKey를 sample_chain로 전달해야 합니다.
Step27:  여러 체인을 실행하기 위해 상태 배치를 sample_chain에 전달하거나 vmap를 사용할 수 있습니다(두 접근법의 성능 차이는 아직 조사하지 않음).
Step28:  옵티마이저
TFP on JAX는 BFGS 및 L-BFGS와 같은 일부 중요한 옵티마이저를 지원합니다. 조정된 간단한 제곱 손실 함수를 설정해보도록 하겠습니다.
Step29:  BFGS는 이러한 손실의 최소값을 찾을 수 있습니다.
Step30:  L-BFGS도 마찬가지입니다.
Step31:  vmap L-BFGS를 하기 위해 단일 시작점에 대한 손실을 최적화하는 함수를 설정해보겠습니다.
Step32:  주의 사항
TF와 JAX 사이에는 몇 가지 근본적인 차이점이 있으며, 일부 TFP 동작은 두 기판 사이에서 다르게 되고 일부 기능을 지원하지 않게 됩니다. 예를 들면, 다음과 같습니다.
TFP on JAX는 tf.Variable와 같은 것을 지원하지 않는데, JAX에는 이러한 것이 존재하지 않기 때문입니다. 이는 또한 tfp.util.TransformedVariable와 같은 유틸리티가 지원되지 않음을 의미합니다.
tfp.layers는 Keras 및 tf.Variable에 대한 종속성으로 인해 아직 백엔드에서 지원되지 않습니다.
tfp.math.minimize는 tf.Variable에 대한 종속성 때문에 TFP on JAX에서 작동하지 않습니다.
TFP on JAX를 사용하면 텐서 형상이 항상 구체적인 정수값이며 TFP on TF에서와 같이 알 수 없거나 동적이지 않습니다.
의사 난수는 TF와 JAX에서 다르게 처리됩니다(부록 참조).
tfp.experimental의 라이브러리는 JAX 기판에 존재할 것으로 보장되지 않습니다.
Dtype 프로모션 규칙은 TF와 JAX 사이에서 다릅니다. TFP on JAX는 일관성을 위해 내부적으로 TF의 dtype 의미 체계를 지키려고 합니다.
Bijectors는 아직 JAX pytrees로 등록되지 않았습니다.
TFP on JAX에서 지원되는 내용의 전체 목록을 보려면 API 설명서를 참조하세요.
결론
우리는 다수의 TFP 기능을 JAX로 이식했으며 모두가 무엇을 구축하게 될지 기대하고 있습니다. 일부 기능은 아직 지원되지 않습니다. 여러분을 위해 우리가 놓친 부분이 있다면(또는 버그를 발견한다면!) 연락해 주시기 바랍니다. tfprobability@tensorflow.org로 이메일을 보내거나 Github 리포지토리에서 문제를 보고할 수 있습니다.
부록
Step33:  JAX에서 난수 함수는 키를 사용하여 결정론적으로 난수 변량을 생성하므로 다시 사용해서는 안 됩니다. 예를 들어 key를 사용하여 정규 분포 값을 샘플링할 수는 있지만 다른 곳에서 key를 다시 사용해서는 안 됩니다. 또한, random.normal에 동일한 값을 전달하면 동일한 값이 생성됩니다.
Step34:  그렇다면 어떻게 해야 단일 키로 여러 샘플을 이끌어 낼 수 있을까요? 정답은 키 분할입니다. 기본 개념은 PRNGKey를 여러 개로 분할할 수 있으며, 각각의 새 키를 독립적인 난수 소스로 취급할 수 있다는 것입니다.
Step35:  키 분할은 결정론적이지만 혼란스럽기 때문에 이제 각각의 새 키를 사용하여 고유한 난수 샘플을 이끌어낼 수 있습니다.
<ASSISTANT_TASK:>
Python Code:
#@title Licensed under the Apache License, Version 2.0 (the "License"); { display-mode: "form" }
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
Explanation: Copyright 2020 The TensorFlow Probability Authors.
Licensed under the Apache License, Version 2.0 (the "License");
End of explanation
!pip uninstall tensorflow -y -q
Explanation: JAX에서 TensorFlow 확률(TFP on JAX)
<table class="tfo-notebook-buttons" align="left">
  <td>     <a target="_blank" href="https://www.tensorflow.org/probability/examples/TensorFlow_Probability_on_JAX"><img src="https://www.tensorflow.org/images/tf_logo_32px.png">TensorFlow.org에서 보기</a>   </td>
  <td><a target="_blank" href="https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/ko/probability/examples/TensorFlow_Probability_on_JAX.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png">Google Colab에서 실행</a></td>
  <td><a target="_blank" href="https://github.com/tensorflow/docs-l10n/blob/master/site/ko/probability/examples/TensorFlow_Probability_on_JAX.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png">GitHub에서 소스 보기</a></td>
  <td><a href="https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/ko/probability/examples/TensorFlow_Probability_on_JAX.ipynb"><img src="https://www.tensorflow.org/images/download_logo_32px.png">노트북 다운로드</a></td>
</table>
TensorFlow 확률(TFP)은 이제 JAX에서도 작동하는 확률적 추론 및 통계 분석을 위한 라이브러리입니다! 익숙하지 않은 분들을 위해 말하자면, JAX는 구성 가능한 함수 변환을 기반으로 가속화된 수치 계산을 수행하기 위한 라이브러리입니다.
TFP on JAX는 많은 TFP 사용자들이 현재 편리하게 사용하는 추상화 및 API를 유지하는 한편 정규 TFP의 가장 유용한 기능을 다수 지원합니다.
설정
TFP on JAX는 TensorFlow에 종속적이지 않습니다. 이 Colab에서 TensorFlow를 완전히 제거하도록 합니다.
End of explanation
!pip install -Uq tfp-nightly[jax] > /dev/null
Explanation: TFP의 최신 야간 빌드를 사용하여 TFP on JAX를 설치할 수 있습니다.
End of explanation
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
from sklearn import datasets
sns.set(style='white')
Explanation: 몇 가지 유용한 Python 라이브러리를 가져옵니다.
End of explanation
import jax.numpy as jnp
from jax import grad
from jax import jit
from jax import random
from jax import value_and_grad
from jax import vmap
Explanation: 또한 몇 가지 기본 JAX 기능을 가져옵니다.
End of explanation
from tensorflow_probability.substrates import jax as tfp
tfd = tfp.distributions
tfb = tfp.bijectors
tfpk = tfp.math.psd_kernels
Explanation: TFP on JAX 가져오기
TFP on JAX를 사용하려면 jax "기판"을 가져온 후 평소의 tfp처럼 사용하면 됩니다.
End of explanation
iris = datasets.load_iris()
features, labels = iris['data'], iris['target']
num_features = features.shape[-1]
num_classes = len(iris.target_names)
Explanation: 데모: 베이지안 로지스틱 회귀
JAX 백엔드로 무엇을 할 수 있는지 보여주기 위해 클래식 Iris 데이터세트에 적용된 베이지안 로지스틱 회귀를 구현하겠습니다.
먼저 Iris 데이터세트를 가져온 후 일부 메타데이터를 추출합니다.
End of explanation
Root = tfd.JointDistributionCoroutine.Root
def model():
  w = yield Root(tfd.Sample(tfd.Normal(0., 1.),
                            sample_shape=(num_features, num_classes)))
  b = yield Root(
      tfd.Sample(tfd.Normal(0., 1.), sample_shape=(num_classes,)))
  logits = jnp.dot(features, w) + b
  yield tfd.Independent(tfd.Categorical(logits=logits),
                        reinterpreted_batch_ndims=1)
dist = tfd.JointDistributionCoroutine(model)
def target_log_prob(*params):
  return dist.log_prob(params + (labels,))
Explanation: tfd.JointDistributionCoroutine를 사용하여 모델을 정의할 수 있습니다. 가중치와 바이어스 항 모두에 표준 정규 사전 분포를 넣은 후 샘플링된 레이블을 데이터에 고정하는 target_log_prob 함수를 작성합니다.
End of explanation
init_key, sample_key = random.split(random.PRNGKey(0))
init_params = tuple(dist.sample(seed=init_key)[:-1])
@jit
def run_chain(key, state):
  kernel = tfp.mcmc.NoUTurnSampler(target_log_prob, 1e-3)
  return tfp.mcmc.sample_chain(500,
      current_state=state,
      kernel=kernel,
      trace_fn=lambda _, results: results.target_log_prob,
      num_burnin_steps=500,
      seed=key)
states, log_probs = run_chain(sample_key, init_params)
plt.figure()
plt.plot(log_probs)
plt.ylabel('Target Log Prob')
plt.xlabel('Iterations of NUTS')
plt.show()
Explanation: dist에서 샘플링하여 MCMC의 초기 상태를 생성합니다. 그런 다음 무작위 키와 초기 상태를 취하는 함수를 정의하고 No-U-Turn-Sampler(NUTS)에서 500개의 샘플을 생성할 수 있습니다. jit과 같은 JAX 변환을 사용하여 XLA로 NUTS 샘플러를 컴파일할 수 있습니다.
End of explanation
def classifier_probs(params):
  dists, _ = dist.sample_distributions(seed=random.PRNGKey(0),
                                       value=params + (None,))
  return dists[-1].distribution.probs_parameter()
Explanation: 샘플을 사용하여 각 가중치 세트의 예측 확률을 평균화하는 방식으로 베이지안 모델 평균화(BMA)를 수행하겠습니다.
먼저 주어진 매개변수 세트에 대해 각 클래스에 대한 확률을 생성하는 함수를 작성합니다. dist.sample_distributions를 사용하여 모델의 최종 분포를 얻을 수 있습니다.
End of explanation
all_probs = jit(vmap(classifier_probs))(states)
print('Average accuracy:', jnp.mean(all_probs.argmax(axis=-1) == labels))
print('BMA accuracy:', jnp.mean(all_probs.mean(axis=0).argmax(axis=-1) == labels))
Explanation: 샘플 세트에 vmap(classifier_probs)를 수행하여 각 샘플의 예측된 클래스 확률을 얻을 수 있습니다. 그런 다음 각 샘플의 평균 정확성과 베이지안 모델 평균화의 정확성을 계산합니다.
End of explanation
dist = tfd.Normal(0., 1.)
print(dist.log_prob(0.))
Explanation: BMA가 오류율을 거의 3분의 1로 줄이는 것으로 보입니다!
기본 사항
TFP on JAX에는 tf.Tensor와 같은 TF 객체를 수락하는 대신 JAX 아날로그를 수락하는 TF와 동일한 API가 있습니다. 예를 들어 tf.Tensor가 이전에 입력으로 사용된 곳이라면 어디에서나 API가 AX DeviceArray를 예상합니다. TFP 메서드는 tf.Tensor를 반환하는 대신에 DeviceArray를 반환합니다. TFP on JAX는 DeviceArray 목록이나 사전과 같은 JAX 객체의 중첩 구조에서도 작동합니다.
분포
대부분의 TFP 분포는 TF 대응 부분과 매우 유사한 의미 체계로 JAX에서 지원됩니다. 또한, JAX Pytrees로 등록되어 JAX 변환 함수의 입력 및 출력이 될 수도 있습니다.
기본 분포
이 분포의 log_prob 메서드는 동일하게 작동합니다.
End of explanation
tfd.Normal(0., 1.).sample(seed=random.PRNGKey(0))
Explanation: 한 분포로부터 샘플링하려면 PRNGKey(또는 정수 목록)에서 seed 키워드 인수로 명시적으로 전달해야 합니다. 명시적으로 시드를 전달하지 못하면 오류가 발생합니다.
End of explanation
dist = tfd.MultivariateNormalDiag(
    loc=jnp.zeros(5),
    scale_diag=jnp.ones(5)
)
print('Event shape:', dist.event_shape)
print('Batch shape:', dist.batch_shape)
Explanation: 분포의 형상 의미 체계는 JAX에서 동일하게 유지됩니다. 여기서 분포는 각각 event_shape와 batch_shape를 갖게 되며 많은 샘플을 이끌어낼수록 sample_shape 차원이 더 많이 추가됩니다.
예를 들어 벡터 매개변수를 가진 tfd.MultivariateNormalDiag는 벡터 이벤트 형상과 빈 배치 형상을 갖게 됩니다.
End of explanation
dist = tfd.Normal(
    loc=jnp.ones(5),
    scale=jnp.ones(5),
)
print('Event shape:', dist.event_shape)
print('Batch shape:', dist.batch_shape)
Explanation: 반면에 벡터로 매개변수화된 tfd.Normal는 스칼라 이벤트 형상과 벡터 배치 형상을 갖게 됩니다.
End of explanation
dist =  tfd.Normal(jnp.zeros(5), jnp.ones(5))
s = dist.sample(sample_shape=(10, 2), seed=random.PRNGKey(0))
print(dist.log_prob(s).shape)
dist =  tfd.Independent(tfd.Normal(jnp.zeros(5), jnp.ones(5)), 1)
s = dist.sample(sample_shape=(10, 2), seed=random.PRNGKey(0))
print(dist.log_prob(s).shape)
Explanation: 샘플의 log_prob를 취하는 의미 체계는 JAX에서도 동일하게 작동합니다.
End of explanation
sns.distplot(tfd.Normal(0., 1.).sample(1000, seed=random.PRNGKey(0)))
plt.show()
Explanation: JAX DeviceArray는 NumPy 및 Matplotlib와 같은 라이브러리와 호환되므로 샘플을 직접 플로팅 함수에 공급할 수 있습니다.
End of explanation
sns.distplot(jit(vmap(lambda key: tfd.Normal(0., 1.).sample(seed=key)))(
    random.split(random.PRNGKey(0), 2000)))
plt.show()
x = jnp.linspace(-5., 5., 100)
plt.plot(x, jit(vmap(grad(tfd.Normal(0., 1.).prob)))(x))
plt.show()
Explanation: Distribution 메서드는 JAX 변환과 호환됩니다.
End of explanation
@jit
def random_distribution(key):
  loc_key, scale_key = random.split(key)
  loc, log_scale = random.normal(loc_key), random.normal(scale_key)
  return tfd.Normal(loc, jnp.exp(log_scale))
random_dist = random_distribution(random.PRNGKey(0))
print(random_dist.mean(), random_dist.variance())
Explanation: TFP 분포는 JAX pytree 노드로 등록되어 있기 때문에 분포를 입력 또는 출력으로 사용하여 함수를 작성하고 이러한 함수를 jit를 사용하여 변환할 수는 있지만 아직 vmap를 적용한 함수에 대한 인수로는 지원되지 않습니다.
End of explanation
dist = tfd.TransformedDistribution(
    tfd.Normal(0., 1.),
    tfb.Sigmoid()
)
sns.distplot(dist.sample(1000, seed=random.PRNGKey(0)))
plt.show()
Explanation: 변환된 분포
샘플이 Bijector를 통해 전달된 분포와 같이 변환된 분포도 기본적으로 작동합니다(bijectors도 작동합니다! 아래 참조).
End of explanation
dist = tfd.JointDistributionSequential([
  tfd.Normal(0., 1.),
  lambda x: tfd.Normal(x, 1e-1)
])
plt.scatter(*dist.sample(1000, seed=random.PRNGKey(0)), alpha=0.5)
plt.show()
joint = tfd.JointDistributionNamed(dict(
    e=             tfd.Exponential(rate=1.),
    n=             tfd.Normal(loc=0., scale=2.),
    m=lambda n, e: tfd.Normal(loc=n, scale=e),
    x=lambda    m: tfd.Sample(tfd.Bernoulli(logits=m), 12),
))
joint.sample(seed=random.PRNGKey(0))
Root = tfd.JointDistributionCoroutine.Root
def model():
  e = yield Root(tfd.Exponential(rate=1.))
  n = yield Root(tfd.Normal(loc=0, scale=2.))
  m = yield tfd.Normal(loc=n, scale=e)
  x = yield tfd.Sample(tfd.Bernoulli(logits=m), 12)
joint = tfd.JointDistributionCoroutine(model)
joint.sample(seed=random.PRNGKey(0))
Explanation: 결합 분포
TFP는 구성 요소 분포를 여러 무작위 변수에 대한 단일 분포로 결합할 수 있도록 JointDistribution를 제공합니다. 현재 TFP는 JAX에서 지원되는 3가지 핵심 변화형(JointDistributionSequential, JointDistributionNamed 및 JointDistributionCoroutine)을 제공합니다. AutoBatched 변화형도 모두 지원합니다.
End of explanation
k1, k2, k3 = random.split(random.PRNGKey(0), 3)
observation_noise_variance = 0.01
f = lambda x: jnp.sin(10*x[..., 0]) * jnp.exp(-x[..., 0]**2)
observation_index_points = random.uniform(
    k1, [50], minval=-1.,maxval= 1.)[..., jnp.newaxis]
observations = f(observation_index_points) + tfd.Normal(
    loc=0., scale=jnp.sqrt(observation_noise_variance)).sample(seed=k2)
index_points = jnp.linspace(-1., 1., 100)[..., jnp.newaxis]
kernel = tfpk.ExponentiatedQuadratic(length_scale=0.1)
gprm = tfd.GaussianProcessRegressionModel(
    kernel=kernel,
    index_points=index_points,
    observation_index_points=observation_index_points,
    observations=observations,
    observation_noise_variance=observation_noise_variance)
samples = gprm.sample(10, seed=k3)
for i in range(10):
  plt.plot(index_points, samples[i], alpha=0.5)
plt.plot(observation_index_points, observations, marker='o', linestyle='')
plt.show()
Explanation: 기타 분포
가우시안 프로세스도 JAX 모드에서 작동합니다!
End of explanation
initial_distribution = tfd.Categorical(probs=[0.8, 0.2])
transition_distribution = tfd.Categorical(probs=[[0.7, 0.3],
                                                 [0.2, 0.8]])
observation_distribution = tfd.Normal(loc=[0., 15.], scale=[5., 10.])
model = tfd.HiddenMarkovModel(
    initial_distribution=initial_distribution,
    transition_distribution=transition_distribution,
    observation_distribution=observation_distribution,
    num_steps=7)
print(model.mean())
print(model.log_prob(jnp.zeros(7)))
print(model.sample(seed=random.PRNGKey(0)))
Explanation: Hidden Markov 모델도 지원합니다.
End of explanation
tfb.Exp().inverse(1.)
bij = tfb.Shift(1.)(tfb.Scale(3.))
print(bij.forward(jnp.ones(5)))
print(bij.inverse(jnp.ones(5)))
b = tfb.FillScaleTriL(diag_bijector=tfb.Exp(), diag_shift=None)
print(b.forward(x=[0., 0., 0.]))
print(b.inverse(y=[[1., 0], [.5, 2]]))
b = tfb.Chain([tfb.Exp(), tfb.Softplus()])
# or:
# b = tfb.Exp()(tfb.Softplus())
print(b.forward(-jnp.ones(5)))
Explanation: TensorFlow 또는 XLA 비호환성에 대한 엄격한 종속성으로 인해 PixelCNN와 같은 몇몇 분포는 아직 지원하지 않습니다.
Bijectors
현재 JAX는 대부분의 TFP Bijector를 지원하고 있습니다!
End of explanation
jit(vmap(tfb.Exp().inverse))(jnp.arange(4.))
x = jnp.linspace(0., 1., 100)
plt.plot(x, jit(grad(lambda x: vmap(tfb.Sigmoid().inverse)(x).sum()))(x))
plt.show()
Explanation: Bijector는 jit, grad 및 vmap과 같은 JAX 변환과 호환됩니다.
End of explanation
target_log_prob = tfd.MultivariateNormalDiag(jnp.zeros(2), jnp.ones(2)).log_prob
Explanation: RealNVP 및 FFJORD와 같은 일부 Bijector는 아직 지원하지 않습니다.
MCMC
tfp.mcmc도 JAX로 이식했으므로 JAX에서 Hamiltonian Monte Carlo(HMC) 및 No-U-Turn-Sampler(NUTS)와 같은 알고리즘도 실행할 수 있습니다.
End of explanation
def run_chain(key, state):
  kernel = tfp.mcmc.NoUTurnSampler(target_log_prob, 1e-1)
  return tfp.mcmc.sample_chain(1000,
      current_state=state,
      kernel=kernel,
      trace_fn=lambda _, results: results.target_log_prob,
      seed=key)
states, log_probs = jit(run_chain)(random.PRNGKey(0), jnp.zeros(2))
plt.figure()
plt.scatter(*states.T, alpha=0.5)
plt.figure()
plt.plot(log_probs)
plt.show()
Explanation: TFP on TF와 달리 seed 키워드 인수를 사용하여 PRNGKey를 sample_chain로 전달해야 합니다.
End of explanation
states, log_probs = jit(run_chain)(random.PRNGKey(0), jnp.zeros([10, 2]))
plt.figure()
for i in range(10):
  plt.scatter(*states[:, i].T, alpha=0.5)
plt.figure()
for i in range(10):
  plt.plot(log_probs[:, i], alpha=0.5)
plt.show()
Explanation: 여러 체인을 실행하기 위해 상태 배치를 sample_chain에 전달하거나 vmap를 사용할 수 있습니다(두 접근법의 성능 차이는 아직 조사하지 않음).
End of explanation
minimum = jnp.array([1.0, 1.0])  # The center of the quadratic bowl.
scales = jnp.array([2.0, 3.0])  # The scales along the two axes.
# The objective function and the gradient.
def quadratic_loss(x):
  return jnp.sum(scales * jnp.square(x - minimum))
start = jnp.array([0.6, 0.8])  # Starting point for the search.
Explanation: 옵티마이저
TFP on JAX는 BFGS 및 L-BFGS와 같은 일부 중요한 옵티마이저를 지원합니다. 조정된 간단한 제곱 손실 함수를 설정해보도록 하겠습니다.
End of explanation
optim_results = tfp.optimizer.bfgs_minimize(
    value_and_grad(quadratic_loss), initial_position=start, tolerance=1e-8)
# Check that the search converged
assert(optim_results.converged)
# Check that the argmin is close to the actual value.
np.testing.assert_allclose(optim_results.position, minimum)
# Print out the total number of function evaluations it took. Should be 5.
print("Function evaluations: %d" % optim_results.num_objective_evaluations)
Explanation: BFGS는 이러한 손실의 최소값을 찾을 수 있습니다.
End of explanation
optim_results = tfp.optimizer.lbfgs_minimize(
    value_and_grad(quadratic_loss), initial_position=start, tolerance=1e-8)
# Check that the search converged
assert(optim_results.converged)
# Check that the argmin is close to the actual value.
np.testing.assert_allclose(optim_results.position, minimum)
# Print out the total number of function evaluations it took. Should be 5.
print("Function evaluations: %d" % optim_results.num_objective_evaluations)
Explanation: L-BFGS도 마찬가지입니다.
End of explanation
def optimize_single(start):
  return tfp.optimizer.lbfgs_minimize(
      value_and_grad(quadratic_loss), initial_position=start, tolerance=1e-8)
all_results = jit(vmap(optimize_single))(
    random.normal(random.PRNGKey(0), (10, 2)))
assert all(all_results.converged)
for i in range(10):
  np.testing.assert_allclose(optim_results.position[i], minimum)
print("Function evaluations: %s" % all_results.num_objective_evaluations)
Explanation: vmap L-BFGS를 하기 위해 단일 시작점에 대한 손실을 최적화하는 함수를 설정해보겠습니다.
End of explanation
key = random.PRNGKey(0)  # Creates a key with value [0, 0]
print(key)
Explanation: 주의 사항
TF와 JAX 사이에는 몇 가지 근본적인 차이점이 있으며, 일부 TFP 동작은 두 기판 사이에서 다르게 되고 일부 기능을 지원하지 않게 됩니다. 예를 들면, 다음과 같습니다.
TFP on JAX는 tf.Variable와 같은 것을 지원하지 않는데, JAX에는 이러한 것이 존재하지 않기 때문입니다. 이는 또한 tfp.util.TransformedVariable와 같은 유틸리티가 지원되지 않음을 의미합니다.
tfp.layers는 Keras 및 tf.Variable에 대한 종속성으로 인해 아직 백엔드에서 지원되지 않습니다.
tfp.math.minimize는 tf.Variable에 대한 종속성 때문에 TFP on JAX에서 작동하지 않습니다.
TFP on JAX를 사용하면 텐서 형상이 항상 구체적인 정수값이며 TFP on TF에서와 같이 알 수 없거나 동적이지 않습니다.
의사 난수는 TF와 JAX에서 다르게 처리됩니다(부록 참조).
tfp.experimental의 라이브러리는 JAX 기판에 존재할 것으로 보장되지 않습니다.
Dtype 프로모션 규칙은 TF와 JAX 사이에서 다릅니다. TFP on JAX는 일관성을 위해 내부적으로 TF의 dtype 의미 체계를 지키려고 합니다.
Bijectors는 아직 JAX pytrees로 등록되지 않았습니다.
TFP on JAX에서 지원되는 내용의 전체 목록을 보려면 API 설명서를 참조하세요.
결론
우리는 다수의 TFP 기능을 JAX로 이식했으며 모두가 무엇을 구축하게 될지 기대하고 있습니다. 일부 기능은 아직 지원되지 않습니다. 여러분을 위해 우리가 놓친 부분이 있다면(또는 버그를 발견한다면!) 연락해 주시기 바랍니다. tfprobability@tensorflow.org로 이메일을 보내거나 Github 리포지토리에서 문제를 보고할 수 있습니다.
부록: JAX의 의사 난수
JAX의 의사 난수 생성(PRNG) 모델은 상태 비저장 모델입니다. 상태 저장 모델과 달리 각 난수 추첨 후에 진화하는 변경 가능한 전역 상태가 없습니다. JAX 모델에서 우리는 PRNG 키로 시작하며, 이는 32비트 정수 쌍으로 작동합니다. jax.random.PRNGKey를 사용하여 이러한 키를 구성할 수 있습니다.
End of explanation
print(random.normal(key))
Explanation: JAX에서 난수 함수는 키를 사용하여 결정론적으로 난수 변량을 생성하므로 다시 사용해서는 안 됩니다. 예를 들어 key를 사용하여 정규 분포 값을 샘플링할 수는 있지만 다른 곳에서 key를 다시 사용해서는 안 됩니다. 또한, random.normal에 동일한 값을 전달하면 동일한 값이 생성됩니다.
End of explanation
key1, key2 = random.split(key, num=2)
print(key1, key2)
Explanation: 그렇다면 어떻게 해야 단일 키로 여러 샘플을 이끌어 낼 수 있을까요? 정답은 키 분할입니다. 기본 개념은 PRNGKey를 여러 개로 분할할 수 있으며, 각각의 새 키를 독립적인 난수 소스로 취급할 수 있다는 것입니다.
End of explanation
print(random.normal(key1), random.normal(key2))
Explanation: 키 분할은 결정론적이지만 혼란스럽기 때문에 이제 각각의 새 키를 사용하여 고유한 난수 샘플을 이끌어낼 수 있습니다.
End of explanation
<END_TASK>