                send_start_params(ws)
                send_audio(ws)
                send_finish(ws)
                logger.debug("thread terminating")

import os
import argparse
import tkinter
from tkinter import Label, Tk, PhotoImage, Entry, LEFT, W, END, Button, E
import time
import threading
import json
import uuid
import yaml
from attrdict import AttrDict
import _locale
import jieba
import paddle
from paddlenlp.data import Vocab
from paddlenlp.transformers import position_encoding_init
from paddlenlp.utils.log import logger
from subword_nmt import subword_nmt
import websocket

open_speech = True
try:
    from pyaudio import PyAudio, paInt16
except ImportError as e:
    open_speech = False
    logger.warning("No module named 'pyaudio', so no audio demo.")

import const
from model_demo import SimultaneousTransformerDemo

# By default, the Windows system opens the file with GBK code,
# and the subword_nmt package does not support setting open encoding,
# so it is set to UTF-8 uniformly.
_locale._getdefaultlocale = (lambda *args: ['en_US', 'utf8'])

is_win = False
if os.name == 'nt':
    is_win = True

class STACLTokenizer:
    """
    Jieba+BPE, and convert tokens to ids.
    """

    def __init__(self, args, is_chinese):
        bpe_parser = subword_nmt.create_apply_bpe_parser()
        bpe_args = bpe_parser.parse_args(args=['-c', args.src_bpe_dict])
        self.bpe = subword_nmt.BPE(bpe_args.codes, bpe_args.merges,
                                   bpe_args.separator, None,
                                   bpe_args.glossaries)
        self.is_chinese = is_chinese

        self.src_vocab = Vocab.load_vocabulary(
            args.src_vocab_fpath,
            bos_token=args.special_token[0],
            eos_token=args.special_token[1],
            unk_token=args.special_token[2])

        self.trg_vocab = Vocab.load_vocabulary(
            args.trg_vocab_fpath,
            bos_token=args.special_token[0],
            eos_token=args.special_token[1],
            unk_token=args.special_token[2])

        args.src_vocab_size = len(self.src_vocab)
        args.trg_vocab_size = len(self.trg_vocab)
        self.args = args

    def tokenize(self, raw_string):
        raw_string = raw_string.strip('\n')
        if not raw_string:
            return raw_string, raw_string
        if self.is_chinese:
            raw_string = ' '.join(jieba.cut(raw_string))
        bpe_str = self.bpe.process_line(raw_string)
        ids = self.src_vocab.to_indices(bpe_str.split())
        return bpe_str.split(), ids

def init_model(args, init_from_params):
    # Define model
    args.init_from_params = init_from_params
    transformer = SimultaneousTransformerDemo(
        args.src_vocab_size, args.trg_vocab_size, args.max_length + 1,
        args.n_layer, args.n_head, args.d_model, args.d_inner_hid, args.dropout,
        args.weight_sharing, args.bos_idx, args.eos_idx, args.waitk)

    # Load the trained model
    assert args.init_from_params, (
        "Please set init_from_params to load the infer model.")

    model_dict = paddle.load(
        os.path.join(args.init_from_params, "transformer.pdparams"))

    # To avoid a longer length than training, reset the size of position
    # encoding to max_length
    model_dict["src_pos_embedding.pos_encoder.weight"] = position_encoding_init(
        args.max_length + 1, args.d_model)
    model_dict["trg_pos_embedding.pos_encoder.weight"] = position_encoding_init(
        args.max_length + 1, args.d_model)

    transformer.load_dict(model_dict)
    return transformer

def post_process_seq(seq, bos_idx, eos_idx, output_bos=False, output_eos=False):
    """
    Post-process the decoded sequence.
    """
    eos_pos = len(seq) - 1
    for i, idx in enumerate(seq):
        if idx == eos_idx:
            eos_pos = i
            break
    seq = [
        idx for idx in seq[:eos_pos + 1]
        if (output_bos or idx != bos_idx) and (output_eos or idx != eos_idx)
    ]
    return seq

def translate(args, tokenizer, tokenized_src, transformers, waitks,
              decoder_max_length, is_last, caches, bos_id, all_result):
    # Set evaluate mode
    for transformer in transformers:
        transformer.eval()

    for idx, (waitk, transformer) in enumerate(zip(waitks, transformers)):
        if len(tokenized_src) < waitk or (waitk == -1 and not is_last):
            continue
        with paddle.no_grad():
            input_src = tokenized_src
            if is_last:
                decoder_max_length[idx] = args.max_out_len
                input_src += [args.eos_idx]
            src_word = paddle.to_tensor(input_src).unsqueeze(axis=0)
            finished_seq, finished_scores, cache = transformer.greedy_search(
                src_word,
                max_len=decoder_max_length[idx],
                waitk=waitk,
                caches=caches[idx],
                bos_id=bos_id[idx])
            caches[idx] = cache
            finished_seq = finished_seq.numpy()
            for beam_idx, beam in enumerate(finished_seq[0]):
                if beam_idx >= args.n_best:
                    break
                id_list = post_process_seq(beam, args.bos_idx, args.eos_idx)
                if len(id_list) == 0:
                    continue
                bos_id[idx] = id_list[-1]
                word_list = tokenizer.trg_vocab.to_tokens(id_list)
                for word in word_list:
                    all_result[idx].append(word)
                res = ' '.join(word_list).replace('@@ ', '')
                logger.debug('[waitk={}] {}'.format(waitk, res))

def cut_line(str, line_len):
    """
    Wrap output
    """
    result = []
    temp = []
    for idx, item in enumerate(str.split()):
        temp.append(item)
        if (idx + 1) % line_len == 0:
            result.append(' '.join(temp))
            temp = []
    if len(temp) != 0:
        result.append(' '.join(temp))
    return '\n'.join(result)

def process(args, tokenizer, transformers, waitks):
    """
    GUI and main waitk program
    :param args:
    :param tokenizer:
    :param transformers:
    :param waitks:
    :return:
    """
    font_align = ('Courier', 20)
    font_label = ('Times', 14)

    if is_win:
        font_align = ('Courier', 15)
        font_label = ('Times', 11)

    window = Tk()

    window.title("Welcome to Simultaneous Translation")
    window.geometry('1200x600')

    logo = PhotoImage(file='images/paddlenlp.png')
    button = Label(window, image=logo, compound='center')
    button.place(x=0, y=0)

    lbl1 = Label(
        window,
        text="Chinese input:",
        fg="green",
        font=font_label,
        anchor=E,
        width=28)
    lbl1.place(x=0, y=60)
    txt = Entry(window, font=font_align)
    txt.place(x=250, y=50, width=800, height=50)

    button_on = Button(window, text='REC', relief='raised', cursor="hand2")
    if open_speech:
        button_on.place(x=1090, y=52)

    s_x, s_y = 0, 130
    x, y = 250, 120

    lbl2_s = Label(
        window,
        text="Jieba+BPE:",
        fg="black",
        font=font_label,
        anchor=E,
        width=28)
    lbl2_s.place(x=s_x, y=s_y)
    lbl2 = Label(
        window, text="", font=font_align, background="pale green", anchor=E)
    lbl2.place(x=x, y=y, width=800, height=50)

    waitnum = '1'
    lbl3_s = Label(
        window,
        text="Simultaneous\nTranslation (wait " + waitnum + "):",
        fg="red",
        font=font_label,
        anchor=E,
        width=28)
    lbl3_s.place(x=s_x, y=s_y + 70)

    lbl3 = Label(window, text="", font=font_align, background="linen")
    lbl3.place(x=x, y=y + 75, width=800, height=50)

    waitnum = '3'
    lbl4_s = Label(
        window,
        text="Simultaneous\nTranslation (wait " + waitnum + "):",
        fg="red",
        font=font_label,
        anchor=E,
        width=28)
    lbl4_s.place(x=s_x, y=s_y + 140)
    lbl4 = Label(window, text="", font=font_align, background="linen")
    lbl4.place(x=x, y=y + 145, width=800, height=50)

    waitnum = '5'
    lbl5_s = Label(
        window,
        text="Simultaneous\nTranslation (wait " + waitnum + "):",
        fg="red",
        font=font_label,
        anchor=E,
        width=28)
    lbl5_s.place(x=s_x, y=s_y + 210)
    lbl5 = Label(window, text="", font=font_align, background="linen")
    lbl5.place(x=x, y=y + 215, width=800, height=50)

    lbl6_s = Label(
        window,
        text="Full Sentence\nTranslation (wait -1):",
        fg="blue",
        font=font_label,
        anchor=E,
        width=28)
    lbl6_s.place(x=s_x, y=s_y + 280)

    lbl6 = Label(window, text="", font=font_align, background="sky blue")
    lbl6.place(x=x, y=y + 285, width=800, height=50)

    def set_val(event=None):
        """
        Start translating
        """
        global i
        global caches
        global bos_id
        global decoder_max_length
        global all_result
        global is_last
        global user_input_bpe
        global user_input_tokenized
        bpe_str, tokenized_src = tokenizer.tokenize(txt.get())
        while i < len(tokenized_src):
            user_input_bpe.append(bpe_str[i])
            user_input_tokenized.append(tokenized_src[i])
            lbl2.configure(
                text=cut_line((lbl2.cget("text") + ' ' + bpe_str[i]).strip(),
                              20),
                fg="black",
                anchor=W,
                justify=LEFT)
            window.update()
            if bpe_str[i] in ['。', '？', '！']:
                is_last = True
            translate(args, tokenizer, user_input_tokenized, transformers,
                      waitks, decoder_max_length, is_last, caches, bos_id,
                      all_result)
            lbl3.configure(
                text=cut_line(' '.join(all_result[0]).replace('@@ ', ''), 11),
                fg="red",
                anchor=W,
                justify=LEFT)
            lbl4.configure(
                text=cut_line(' '.join(all_result[1]).replace('@@ ', ''), 11),
                fg="red",
                anchor=W,
                justify=LEFT)
            lbl5.configure(
                text=cut_line(' '.join(all_result[2]).replace('@@ ', ''), 11),
                fg="red",
                anchor=W,
                justify=LEFT)
            lbl6.configure(
                text=cut_line(' '.join(all_result[3]).replace('@@ ', ''), 11),
                fg="blue",
                anchor=W,
                justify=LEFT)
            window.update()
            if is_last:
                caches = [None] * len(waitks)
                bos_id = [None] * len(waitks)
                decoder_max_length = [1] * len(waitks)
                is_last = False
                user_input_bpe = []
                user_input_tokenized = []
            i += 1

    def set_val_voice(event=None):
        """
        Start translating
        """

        def send_start_params(ws):
            """
            Send start frame
            :param websocket.WebSocket ws:
            :return:
            """
            req = {
                "type": "START",
                "data": {
                    "appid": const.APPID,
                    "appkey": const.APPKEY,
                    "dev_pid": const.DEV_PID,
                    "cuid": "yourself_defined_user_id",
                    "sample": 16000,
                    "format": "pcm"
                }
            }
            body = json.dumps(req)
            ws.send(body, websocket.ABNF.OPCODE_TEXT)
            logger.info("send START frame with params:" + body)

        def send_audio(ws):
            """
             Send audio
            :param  websocket.WebSocket ws:
            :return:
            """
            # 160ms record
            chunk_ms = 160

            chunk_len = int(16000 * 2 / 1000 * chunk_ms)

            pa = PyAudio()
            stream = pa.open(
                format=paInt16,
                channels=1,
                rate=16000,
                input=True,
                frames_per_buffer=chunk_len // 2)

            while True:
                frames = []
                frame = stream.read(chunk_len // 2, exception_on_overflow=False)
                frames.append(frame)
                body = b''.join(frames)
                if len(body) == 0:
                    logger.info("empty body")
                    continue
                logger.debug("try to send audio length {}".format(len(body)))
                ws.send(body, websocket.ABNF.OPCODE_BINARY)

        def send_finish(ws):
            """
            Send finished frame
            :param websocket.WebSocket ws:
            :return:
            """
            req = {"type": "FINISH"}
            body = json.dumps(req)
            ws.send(body, websocket.ABNF.OPCODE_TEXT)
            logger.info("send FINISH frame")

        def close_websocket(ws_app):
            if ws_app:
                logger.info('close ws_app.')
                send_finish(ws_app)
                ws_app.close()
            logger.info('ws_app closed.')

        def on_open(ws):
            """
            Send data frame after connected
            :param  websocket.WebSocket ws:
            :return:
            """

            def run(*args):
                """
                Send data frame
                :param args:
                :return:
                """

            threading.Thread(target=run).start()

        def on_error(ws, error):
            """
            For error
            :param ws:
            :param error: json
            :return:
                """
            logger.error("error: " + str(error))

        def on_close(ws):
            """
            Close websocket
            :param websocket.WebSocket ws:
            :return:
            """
            logger.info("ws close ...")

        def on_message(ws, message):
            """
            Response from server
            :param ws:
            :param message: json
            :return:
            """
            global i
            global text
            global caches
            global bos_id
            global decoder_max_length
            global all_result
            global is_last
            global user_input_bpe
            global user_input_tokenized
            global ws_app
            global start_time

            logger.info("Response: " + message)
            message = json.loads(message)
            if is_last and ws_app:
                close_websocket(ws_app)
            end_time = time.time()
            if end_time - start_time > 10 and ws_app:
                close_websocket(ws_app)
                logger.info('ws_app started at: {} closed at: {}, cost {}s.'.
                            format(start_time, end_time, end_time - start_time))
            if 'result' in message:
                start_time = time.time()
                text = message['result']
                txt.delete(0, END)
                txt.insert(0, text)
                bpe_str, tokenized_src = tokenizer.tokenize(txt.get())
                while i < len(tokenized_src):
                    user_input_bpe.append(bpe_str[i])
                    user_input_tokenized.append(tokenized_src[i])
                    lbl2.configure(
                        text=cut_line(
                            (lbl2.cget("text") + ' ' + bpe_str[i]).strip(), 20),
                        fg="black",
                        anchor=W,
                        justify=LEFT)
                    window.update()
                    if bpe_str[i] in ['。', '？', '！']:
                        is_last = True
                    translate(args, tokenizer, user_input_tokenized,
                              transformers, waitks, decoder_max_length, is_last,
                              caches, bos_id, all_result)
                    lbl3.configure(
                        text=cut_line(' '.join(all_result[0]).replace('@@ ',
                                                                      ''), 11),
                        fg="red",
                        anchor=W,
                        justify=LEFT)
                    lbl4.configure(
                        text=cut_line(' '.join(all_result[1]).replace('@@ ',
                                                                      ''), 11),
                        fg="red",
                        anchor=W,
                        justify=LEFT)
                    lbl5.configure(
                        text=cut_line(' '.join(all_result[2]).replace('@@ ',
                                                                      ''), 11),
                        fg="red",
                        anchor=W,
                        justify=LEFT)
                    lbl6.configure(
                        text=cut_line(' '.join(all_result[3]).replace('@@ ',
                                                                      ''), 11),
                        fg="blue",
                        anchor=W,
                        justify=LEFT)
                    window.update()
                    if is_last:
                        caches = [None] * len(waitks)
                        bos_id = [None] * len(waitks)
                        decoder_max_length = [1] * len(waitks)
                        is_last = False
                        user_input_bpe = []
                        user_input_tokenized = []
                        if ws_app:
                            close_websocket(ws_app)
                    i += 1

        logger.info("begin")
        uri = const.URI + "?sn=" + str(uuid.uuid1())
        logger.info("uri is " + uri)
        global start_time
        start_time = time.time()
        global ws_app
        ws_app = websocket.WebSocketApp(
            uri,
            on_open=on_open,
            on_message=on_message,
            on_error=on_error,
            on_close=on_close)
        ws_app.run_forever()

    def clear():
        """
        Clear input and output
        """
        txt.delete(0, END)
        global i
        global text
        global caches
        global bos_id
        global decoder_max_length
        global all_result
        global is_last
        global user_input_bpe
        global user_input_tokenized
        global ws_app
        global start_time
        if ws_app:
            ws_app.close()
        decoder_max_length = [1] * len(waitks)
        caches = [None] * len(waitks)
        bos_id = [None] * len(waitks)
        all_result = [[], [], [], []]
        i = 0
        is_last = False
        user_input_bpe = []
        user_input_tokenized = []
        start_time = 0
        logger.info('CLEAR')
        logger.info(f'i: {i}')
        logger.info(f'caches: {caches}')
        logger.info(f'bos_id: {bos_id}')
        logger.info(f'decoder_max_length: {decoder_max_length}')
        logger.info(f'all_result: {all_result}')
        logger.info(f'is_last: {is_last}')
        lbl2.configure(text="", fg="black", anchor=W, justify=LEFT)
        lbl3.configure(text="", fg="red", anchor=W, justify=LEFT)
        lbl4.configure(text="", fg="red", anchor=W, justify=LEFT)
        lbl5.configure(text="", fg="red", anchor=W, justify=LEFT)
        lbl6.configure(text="", fg="blue", anchor=W, justify=LEFT)
        window.update()

    txt.bind('<Return>', set_val)
    button_on.bind('<Button-1>', set_val_voice)

    desc1 = Label(
        window,
        text='使用说明：1. 在Chinese input输入中文，按【回车键】开始实时翻译，'
        '遇到【。！？】结束整句，按【CLEAR】清空所有的输入和输出；',
        anchor=E)
    desc1.place(x=s_x + 100, y=s_y + 380)

    backspace_cnt = 19
    if is_win:
        backspace_cnt = 15

    desc2 = Label(
        window,
        text=' ' * backspace_cnt + '2. 按【REC】开始录音并开始实时翻译，遇到【。！？】结束整句，'
        '按【CLEAR】清空所有的输入和输出。',
        anchor=E)
    if open_speech:
        desc2.place(x=s_x + 100, y=s_y + 410)

    button_clear = Button(
        window, text='CLEAR', relief="raised", cursor="hand2", command=clear)

    button_clear.place(x=x + 840, y=y + 380)

    window.mainloop()

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--config",
        default="./transformer_demo.yaml",
        type=str,
        help="Path of the config file. ")
    args = parser.parse_args()
    return args

if __name__ == "__main__":
    args = parse_args()
    yaml_file = args.config
    with open(yaml_file, 'rt') as f:
        args = AttrDict(yaml.safe_load(f))

    if args.device == 'gpu':
        place = "gpu:0"
    elif args.device == 'xpu':
        place = "xpu:0"
    elif args.device == 'cpu':
        place = "cpu"
    paddle.set_device(place)

    tokenizer = STACLTokenizer(args, is_chinese=True)
    waitks = [1, 3, 5, -1]

    transformers = []
    for waitk in waitks:
        transformers.append(init_model(args, f'models/nist_wait_{waitk}'))
        logger.info(f'Loaded wait_{waitk} model.')

    # for decoding max length
    decoder_max_length = [1] * len(waitks)
    # for decoding cache
    caches = [None] * len(waitks)
    # for decoding start token id
    bos_id = [None] * len(waitks)
    all_result = [[], [], [], []]
    # current source word index
    i = 0
    is_last = False
    # subword after bpe
    user_input_bpe = []
    # tokenized id
    user_input_tokenized = []
    text = ''
    # websocket app
    ws_app = None
    # start time
    start_time = 0

    process(args, tokenizer, transformers, waitks)


while True:
                frames = []
                frame = stream.read(chunk_len // 2, exception_on_overflow=False)
                frames.append(frame)
                body = b''.join(frames)
                if len(body) == 0:
                    logger.info("empty body")
                    continue
                logger.debug("try to send audio length {}".format(len(body)))
                ws.send(body, websocket.ABNF.OPCODE_BINARY)

import os
import argparse
import tkinter
from tkinter import Label, Tk, PhotoImage, Entry, LEFT, W, END, Button, E
import time
import threading
import json
import uuid
import yaml
from attrdict import AttrDict
import _locale
import jieba
import paddle
from paddlenlp.data import Vocab
from paddlenlp.transformers import position_encoding_init
from paddlenlp.utils.log import logger
from subword_nmt import subword_nmt
import websocket

open_speech = True
try:
    from pyaudio import PyAudio, paInt16
except ImportError as e:
    open_speech = False
    logger.warning("No module named 'pyaudio', so no audio demo.")

import const
from model_demo import SimultaneousTransformerDemo

# By default, the Windows system opens the file with GBK code,
# and the subword_nmt package does not support setting open encoding,
# so it is set to UTF-8 uniformly.
_locale._getdefaultlocale = (lambda *args: ['en_US', 'utf8'])

is_win = False
if os.name == 'nt':
    is_win = True

class STACLTokenizer:
    """
    Jieba+BPE, and convert tokens to ids.
    """

    def __init__(self, args, is_chinese):
        bpe_parser = subword_nmt.create_apply_bpe_parser()
        bpe_args = bpe_parser.parse_args(args=['-c', args.src_bpe_dict])
        self.bpe = subword_nmt.BPE(bpe_args.codes, bpe_args.merges,
                                   bpe_args.separator, None,
                                   bpe_args.glossaries)
        self.is_chinese = is_chinese

        self.src_vocab = Vocab.load_vocabulary(
            args.src_vocab_fpath,
            bos_token=args.special_token[0],
            eos_token=args.special_token[1],
            unk_token=args.special_token[2])

        self.trg_vocab = Vocab.load_vocabulary(
            args.trg_vocab_fpath,
            bos_token=args.special_token[0],
            eos_token=args.special_token[1],
            unk_token=args.special_token[2])

        args.src_vocab_size = len(self.src_vocab)
        args.trg_vocab_size = len(self.trg_vocab)
        self.args = args

    def tokenize(self, raw_string):
        raw_string = raw_string.strip('\n')
        if not raw_string:
            return raw_string, raw_string
        if self.is_chinese:
            raw_string = ' '.join(jieba.cut(raw_string))
        bpe_str = self.bpe.process_line(raw_string)
        ids = self.src_vocab.to_indices(bpe_str.split())
        return bpe_str.split(), ids

def init_model(args, init_from_params):
    # Define model
    args.init_from_params = init_from_params
    transformer = SimultaneousTransformerDemo(
        args.src_vocab_size, args.trg_vocab_size, args.max_length + 1,
        args.n_layer, args.n_head, args.d_model, args.d_inner_hid, args.dropout,
        args.weight_sharing, args.bos_idx, args.eos_idx, args.waitk)

    # Load the trained model
    assert args.init_from_params, (
        "Please set init_from_params to load the infer model.")

    model_dict = paddle.load(
        os.path.join(args.init_from_params, "transformer.pdparams"))

    # To avoid a longer length than training, reset the size of position
    # encoding to max_length
    model_dict["src_pos_embedding.pos_encoder.weight"] = position_encoding_init(
        args.max_length + 1, args.d_model)
    model_dict["trg_pos_embedding.pos_encoder.weight"] = position_encoding_init(
        args.max_length + 1, args.d_model)

    transformer.load_dict(model_dict)
    return transformer

def post_process_seq(seq, bos_idx, eos_idx, output_bos=False, output_eos=False):
    """
    Post-process the decoded sequence.
    """
    eos_pos = len(seq) - 1
    for i, idx in enumerate(seq):
        if idx == eos_idx:
            eos_pos = i
            break
    seq = [
        idx for idx in seq[:eos_pos + 1]
        if (output_bos or idx != bos_idx) and (output_eos or idx != eos_idx)
    ]
    return seq

def translate(args, tokenizer, tokenized_src, transformers, waitks,
              decoder_max_length, is_last, caches, bos_id, all_result):
    # Set evaluate mode
    for transformer in transformers:
        transformer.eval()

    for idx, (waitk, transformer) in enumerate(zip(waitks, transformers)):
        if len(tokenized_src) < waitk or (waitk == -1 and not is_last):
            continue
        with paddle.no_grad():
            input_src = tokenized_src
            if is_last:
                decoder_max_length[idx] = args.max_out_len
                input_src += [args.eos_idx]
            src_word = paddle.to_tensor(input_src).unsqueeze(axis=0)
            finished_seq, finished_scores, cache = transformer.greedy_search(
                src_word,
                max_len=decoder_max_length[idx],
                waitk=waitk,
                caches=caches[idx],
                bos_id=bos_id[idx])
            caches[idx] = cache
            finished_seq = finished_seq.numpy()
            for beam_idx, beam in enumerate(finished_seq[0]):
                if beam_idx >= args.n_best:
                    break
                id_list = post_process_seq(beam, args.bos_idx, args.eos_idx)
                if len(id_list) == 0:
                    continue
                bos_id[idx] = id_list[-1]
                word_list = tokenizer.trg_vocab.to_tokens(id_list)
                for word in word_list:
                    all_result[idx].append(word)
                res = ' '.join(word_list).replace('@@ ', '')
                logger.debug('[waitk={}] {}'.format(waitk, res))

def cut_line(str, line_len):
    """
    Wrap output
    """
    result = []
    temp = []
    for idx, item in enumerate(str.split()):
        temp.append(item)
        if (idx + 1) % line_len == 0:
            result.append(' '.join(temp))
            temp = []
    if len(temp) != 0:
        result.append(' '.join(temp))
    return '\n'.join(result)

def process(args, tokenizer, transformers, waitks):
    """
    GUI and main waitk program
    :param args:
    :param tokenizer:
    :param transformers:
    :param waitks:
    :return:
    """
    font_align = ('Courier', 20)
    font_label = ('Times', 14)

    if is_win:
        font_align = ('Courier', 15)
        font_label = ('Times', 11)

    window = Tk()

    window.title("Welcome to Simultaneous Translation")
    window.geometry('1200x600')

    logo = PhotoImage(file='images/paddlenlp.png')
    button = Label(window, image=logo, compound='center')
    button.place(x=0, y=0)

    lbl1 = Label(
        window,
        text="Chinese input:",
        fg="green",
        font=font_label,
        anchor=E,
        width=28)
    lbl1.place(x=0, y=60)
    txt = Entry(window, font=font_align)
    txt.place(x=250, y=50, width=800, height=50)

    button_on = Button(window, text='REC', relief='raised', cursor="hand2")
    if open_speech:
        button_on.place(x=1090, y=52)

    s_x, s_y = 0, 130
    x, y = 250, 120

    lbl2_s = Label(
        window,
        text="Jieba+BPE:",
        fg="black",
        font=font_label,
        anchor=E,
        width=28)
    lbl2_s.place(x=s_x, y=s_y)
    lbl2 = Label(
        window, text="", font=font_align, background="pale green", anchor=E)
    lbl2.place(x=x, y=y, width=800, height=50)

    waitnum = '1'
    lbl3_s = Label(
        window,
        text="Simultaneous\nTranslation (wait " + waitnum + "):",
        fg="red",
        font=font_label,
        anchor=E,
        width=28)
    lbl3_s.place(x=s_x, y=s_y + 70)

    lbl3 = Label(window, text="", font=font_align, background="linen")
    lbl3.place(x=x, y=y + 75, width=800, height=50)

    waitnum = '3'
    lbl4_s = Label(
        window,
        text="Simultaneous\nTranslation (wait " + waitnum + "):",
        fg="red",
        font=font_label,
        anchor=E,
        width=28)
    lbl4_s.place(x=s_x, y=s_y + 140)
    lbl4 = Label(window, text="", font=font_align, background="linen")
    lbl4.place(x=x, y=y + 145, width=800, height=50)

    waitnum = '5'
    lbl5_s = Label(
        window,
        text="Simultaneous\nTranslation (wait " + waitnum + "):",
        fg="red",
        font=font_label,
        anchor=E,
        width=28)
    lbl5_s.place(x=s_x, y=s_y + 210)
    lbl5 = Label(window, text="", font=font_align, background="linen")
    lbl5.place(x=x, y=y + 215, width=800, height=50)

    lbl6_s = Label(
        window,
        text="Full Sentence\nTranslation (wait -1):",
        fg="blue",
        font=font_label,
        anchor=E,
        width=28)
    lbl6_s.place(x=s_x, y=s_y + 280)

    lbl6 = Label(window, text="", font=font_align, background="sky blue")
    lbl6.place(x=x, y=y + 285, width=800, height=50)

    def set_val(event=None):
        """
        Start translating
        """
        global i
        global caches
        global bos_id
        global decoder_max_length
        global all_result
        global is_last
        global user_input_bpe
        global user_input_tokenized
        bpe_str, tokenized_src = tokenizer.tokenize(txt.get())
        while i < len(tokenized_src):
            user_input_bpe.append(bpe_str[i])
            user_input_tokenized.append(tokenized_src[i])
            lbl2.configure(
                text=cut_line((lbl2.cget("text") + ' ' + bpe_str[i]).strip(),
                              20),
                fg="black",
                anchor=W,
                justify=LEFT)
            window.update()
            if bpe_str[i] in ['。', '？', '！']:
                is_last = True
            translate(args, tokenizer, user_input_tokenized, transformers,
                      waitks, decoder_max_length, is_last, caches, bos_id,
                      all_result)
            lbl3.configure(
                text=cut_line(' '.join(all_result[0]).replace('@@ ', ''), 11),
                fg="red",
                anchor=W,
                justify=LEFT)
            lbl4.configure(
                text=cut_line(' '.join(all_result[1]).replace('@@ ', ''), 11),
                fg="red",
                anchor=W,
                justify=LEFT)
            lbl5.configure(
                text=cut_line(' '.join(all_result[2]).replace('@@ ', ''), 11),
                fg="red",
                anchor=W,
                justify=LEFT)
            lbl6.configure(
                text=cut_line(' '.join(all_result[3]).replace('@@ ', ''), 11),
                fg="blue",
                anchor=W,
                justify=LEFT)
            window.update()
            if is_last:
                caches = [None] * len(waitks)
                bos_id = [None] * len(waitks)
                decoder_max_length = [1] * len(waitks)
                is_last = False
                user_input_bpe = []
                user_input_tokenized = []
            i += 1

    def set_val_voice(event=None):
        """
        Start translating
        """

        def send_start_params(ws):
            """
            Send start frame
            :param websocket.WebSocket ws:
            :return:
            """
            req = {
                "type": "START",
                "data": {
                    "appid": const.APPID,
                    "appkey": const.APPKEY,
                    "dev_pid": const.DEV_PID,
                    "cuid": "yourself_defined_user_id",
                    "sample": 16000,
                    "format": "pcm"
                }
            }
            body = json.dumps(req)
            ws.send(body, websocket.ABNF.OPCODE_TEXT)
            logger.info("send START frame with params:" + body)

        def send_audio(ws):
            """
             Send audio
            :param  websocket.WebSocket ws:
            :return:
            """
            # 160ms record
            chunk_ms = 160

            chunk_len = int(16000 * 2 / 1000 * chunk_ms)

            pa = PyAudio()
            stream = pa.open(
                format=paInt16,
                channels=1,
                rate=16000,
                input=True,
                frames_per_buffer=chunk_len // 2)

        def send_finish(ws):
            """
            Send finished frame
            :param websocket.WebSocket ws:
            :return:
            """
            req = {"type": "FINISH"}
            body = json.dumps(req)
            ws.send(body, websocket.ABNF.OPCODE_TEXT)
            logger.info("send FINISH frame")

        def close_websocket(ws_app):
            if ws_app:
                logger.info('close ws_app.')
                send_finish(ws_app)
                ws_app.close()
            logger.info('ws_app closed.')

        def on_open(ws):
            """
            Send data frame after connected
            :param  websocket.WebSocket ws:
            :return:
            """

            def run(*args):
                """
                Send data frame
                :param args:
                :return:
                """
                send_start_params(ws)
                send_audio(ws)
                send_finish(ws)
                logger.debug("thread terminating")

            threading.Thread(target=run).start()

        def on_error(ws, error):
            """
            For error
            :param ws:
            :param error: json
            :return:
                """
            logger.error("error: " + str(error))

        def on_close(ws):
            """
            Close websocket
            :param websocket.WebSocket ws:
            :return:
            """
            logger.info("ws close ...")

        def on_message(ws, message):
            """
            Response from server
            :param ws:
            :param message: json
            :return:
            """
            global i
            global text
            global caches
            global bos_id
            global decoder_max_length
            global all_result
            global is_last
            global user_input_bpe
            global user_input_tokenized
            global ws_app
            global start_time

            logger.info("Response: " + message)
            message = json.loads(message)
            if is_last and ws_app:
                close_websocket(ws_app)
            end_time = time.time()
            if end_time - start_time > 10 and ws_app:
                close_websocket(ws_app)
                logger.info('ws_app started at: {} closed at: {}, cost {}s.'.
                            format(start_time, end_time, end_time - start_time))
            if 'result' in message:
                start_time = time.time()
                text = message['result']
                txt.delete(0, END)
                txt.insert(0, text)
                bpe_str, tokenized_src = tokenizer.tokenize(txt.get())
                while i < len(tokenized_src):
                    user_input_bpe.append(bpe_str[i])
                    user_input_tokenized.append(tokenized_src[i])
                    lbl2.configure(
                        text=cut_line(
                            (lbl2.cget("text") + ' ' + bpe_str[i]).strip(), 20),
                        fg="black",
                        anchor=W,
                        justify=LEFT)
                    window.update()
                    if bpe_str[i] in ['。', '？', '！']:
                        is_last = True
                    translate(args, tokenizer, user_input_tokenized,
                              transformers, waitks, decoder_max_length, is_last,
                              caches, bos_id, all_result)
                    lbl3.configure(
                        text=cut_line(' '.join(all_result[0]).replace('@@ ',
                                                                      ''), 11),
                        fg="red",
                        anchor=W,
                        justify=LEFT)
                    lbl4.configure(
                        text=cut_line(' '.join(all_result[1]).replace('@@ ',
                                                                      ''), 11),
                        fg="red",
                        anchor=W,
                        justify=LEFT)
                    lbl5.configure(
                        text=cut_line(' '.join(all_result[2]).replace('@@ ',
                                                                      ''), 11),
                        fg="red",
                        anchor=W,
                        justify=LEFT)
                    lbl6.configure(
                        text=cut_line(' '.join(all_result[3]).replace('@@ ',
                                                                      ''), 11),
                        fg="blue",
                        anchor=W,
                        justify=LEFT)
                    window.update()
                    if is_last:
                        caches = [None] * len(waitks)
                        bos_id = [None] * len(waitks)
                        decoder_max_length = [1] * len(waitks)
                        is_last = False
                        user_input_bpe = []
                        user_input_tokenized = []
                        if ws_app:
                            close_websocket(ws_app)
                    i += 1

        logger.info("begin")
        uri = const.URI + "?sn=" + str(uuid.uuid1())
        logger.info("uri is " + uri)
        global start_time
        start_time = time.time()
        global ws_app
        ws_app = websocket.WebSocketApp(
            uri,
            on_open=on_open,
            on_message=on_message,
            on_error=on_error,
            on_close=on_close)
        ws_app.run_forever()

    def clear():
        """
        Clear input and output
        """
        txt.delete(0, END)
        global i
        global text
        global caches
        global bos_id
        global decoder_max_length
        global all_result
        global is_last
        global user_input_bpe
        global user_input_tokenized
        global ws_app
        global start_time
        if ws_app:
            ws_app.close()
        decoder_max_length = [1] * len(waitks)
        caches = [None] * len(waitks)
        bos_id = [None] * len(waitks)
        all_result = [[], [], [], []]
        i = 0
        is_last = False
        user_input_bpe = []
        user_input_tokenized = []
        start_time = 0
        logger.info('CLEAR')
        logger.info(f'i: {i}')
        logger.info(f'caches: {caches}')
        logger.info(f'bos_id: {bos_id}')
        logger.info(f'decoder_max_length: {decoder_max_length}')
        logger.info(f'all_result: {all_result}')
        logger.info(f'is_last: {is_last}')
        lbl2.configure(text="", fg="black", anchor=W, justify=LEFT)
        lbl3.configure(text="", fg="red", anchor=W, justify=LEFT)
        lbl4.configure(text="", fg="red", anchor=W, justify=LEFT)
        lbl5.configure(text="", fg="red", anchor=W, justify=LEFT)
        lbl6.configure(text="", fg="blue", anchor=W, justify=LEFT)
        window.update()

    txt.bind('<Return>', set_val)
    button_on.bind('<Button-1>', set_val_voice)

    desc1 = Label(
        window,
        text='使用说明：1. 在Chinese input输入中文，按【回车键】开始实时翻译，'
        '遇到【。！？】结束整句，按【CLEAR】清空所有的输入和输出；',
        anchor=E)
    desc1.place(x=s_x + 100, y=s_y + 380)

    backspace_cnt = 19
    if is_win:
        backspace_cnt = 15

    desc2 = Label(
        window,
        text=' ' * backspace_cnt + '2. 按【REC】开始录音并开始实时翻译，遇到【。！？】结束整句，'
        '按【CLEAR】清空所有的输入和输出。',
        anchor=E)
    if open_speech:
        desc2.place(x=s_x + 100, y=s_y + 410)

    button_clear = Button(
        window, text='CLEAR', relief="raised", cursor="hand2", command=clear)

    button_clear.place(x=x + 840, y=y + 380)

    window.mainloop()

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--config",
        default="./transformer_demo.yaml",
        type=str,
        help="Path of the config file. ")
    args = parser.parse_args()
    return args

if __name__ == "__main__":
    args = parse_args()
    yaml_file = args.config
    with open(yaml_file, 'rt') as f:
        args = AttrDict(yaml.safe_load(f))

    if args.device == 'gpu':
        place = "gpu:0"
    elif args.device == 'xpu':
        place = "xpu:0"
    elif args.device == 'cpu':
        place = "cpu"
    paddle.set_device(place)

    tokenizer = STACLTokenizer(args, is_chinese=True)
    waitks = [1, 3, 5, -1]

    transformers = []
    for waitk in waitks:
        transformers.append(init_model(args, f'models/nist_wait_{waitk}'))
        logger.info(f'Loaded wait_{waitk} model.')

    # for decoding max length
    decoder_max_length = [1] * len(waitks)
    # for decoding cache
    caches = [None] * len(waitks)
    # for decoding start token id
    bos_id = [None] * len(waitks)
    all_result = [[], [], [], []]
    # current source word index
    i = 0
    is_last = False
    # subword after bpe
    user_input_bpe = []
    # tokenized id
    user_input_tokenized = []
    text = ''
    # websocket app
    ws_app = None
    # start time
    start_time = 0

    process(args, tokenizer, transformers, waitks)


    if k > len(mapping):
        raise ValueError("sample larger than population")
    sample = set()
    while len(sample) < k:
        sample.add(weighted_choice(mapping, seed))
    return list(sample)

#    Pieter Swart <example@nnthink.com>
#    All rights reserved.
# Authors: Aric Hagberg (hagberg@lanl.gov)
#          Dan Schult (dschult@colgate.edu)
#          Ben Edwards (example@nnthink.com)
"""
Utilities for generating random numbers, random sequences, and
random selections.
"""

import random
import sys
import networkx as nx
from networkx.utils import py_random_state

# The same helpers for choosing random sequences from distributions
# uses Python's random module

@py_random_state(2)
def powerlaw_sequence(n, exponent=2.0, seed=None):
    """
    Return sample sequence of length n from a power law distribution.
    """
    return [seed.paretovariate(exponent - 1) for i in range(n)]

@py_random_state(2)
def zipf_rv(alpha, xmin=1, seed=None):
    r"""Returns a random value chosen from the Zipf distribution.

    The return value is an integer drawn from the probability distribution

    .. math::

        p(x)=\frac{x^{-\alpha}}{\zeta(\alpha, x_{\min})},

    where $\zeta(\alpha, x_{\min})$ is the Hurwitz zeta function.

    Parameters
    ----------
    alpha : float
      Exponent value of the distribution
    xmin : int
      Minimum value
    seed : integer, random_state, or None (default)
        Indicator of random number generation state.
        See :ref:`Randomness<randomness>`.

    Returns
    -------
    x : int
      Random value from Zipf distribution

    Raises
    ------
    ValueError:
      If xmin < 1 or
      If alpha <= 1

    Notes
    -----
    The rejection algorithm generates random values for a the power-law
    distribution in uniformly bounded expected time dependent on
    parameters.  See [1]_ for details on its operation.

    Examples
    --------
    >>> nx.zipf_rv(alpha=2, xmin=3, seed=42)  # doctest: +SKIP

    References
    ----------
    .. [1] Luc Devroye, Non-Uniform Random Variate Generation,
       Springer-Verlag, New York, 1986.
    """
    if xmin < 1:
        raise ValueError("xmin < 1")
    if alpha <= 1:
        raise ValueError("a <= 1.0")
    a1 = alpha - 1.0
    b = 2**a1
    while True:
        u = 1.0 - seed.random()  # u in (0,1]
        v = seed.random()  # v in [0,1)
        x = int(xmin * u**-(1.0 / a1))
        t = (1.0 + (1.0 / x))**a1
        if v * x * (t - 1.0) / (b - 1.0) <= t / b:
            break
    return x

def cumulative_distribution(distribution):
    """Returns normalized cumulative distribution from discrete distribution."""

    cdf = [0.0]
    psum = float(sum(distribution))
    for i in range(0, len(distribution)):
        cdf.append(cdf[i] + distribution[i] / psum)
    return cdf

@py_random_state(3)
def discrete_sequence(n, distribution=None, cdistribution=None, seed=None):
    """
    Return sample sequence of length n from a given discrete distribution
    or discrete cumulative distribution.

    One of the following must be specified.

    distribution = histogram of values, will be normalized

    cdistribution = normalized discrete cumulative distribution

    """
    import bisect

    if cdistribution is not None:
        cdf = cdistribution
    elif distribution is not None:
        cdf = cumulative_distribution(distribution)
    else:
        raise nx.NetworkXError(
            "discrete_sequence: distribution or cdistribution missing")

    # get a uniform random number
    inputseq = [seed.random() for i in range(n)]

    # choose from CDF
    seq = [bisect.bisect_left(cdf, s) - 1 for s in inputseq]
    return seq

@py_random_state(2)
def random_weighted_sample(mapping, k, seed=None):
    """Returns k items without replacement from a weighted sample.

    The input is a dictionary of items with weights as values.
    """

@py_random_state(1)
def weighted_choice(mapping, seed=None):
    """Returns a single element from a weighted sample.

    The input is a dictionary of items with weights as values.
    """
    # use roulette method
    rnd = seed.random() * sum(mapping.values())
    for k, w in mapping.items():
        rnd -= w
        if rnd < 0:
            return k


while len(sample) < k:
        sample.add(weighted_choice(mapping, seed))

#    Pieter Swart <example@nnthink.com>
#    All rights reserved.
# Authors: Aric Hagberg (hagberg@lanl.gov)
#          Dan Schult (dschult@colgate.edu)
#          Ben Edwards (example@nnthink.com)
"""
Utilities for generating random numbers, random sequences, and
random selections.
"""

import random
import sys
import networkx as nx
from networkx.utils import py_random_state

# The same helpers for choosing random sequences from distributions
# uses Python's random module

@py_random_state(2)
def powerlaw_sequence(n, exponent=2.0, seed=None):
    """
    Return sample sequence of length n from a power law distribution.
    """
    return [seed.paretovariate(exponent - 1) for i in range(n)]

@py_random_state(2)
def zipf_rv(alpha, xmin=1, seed=None):
    r"""Returns a random value chosen from the Zipf distribution.

    The return value is an integer drawn from the probability distribution

    .. math::

        p(x)=\frac{x^{-\alpha}}{\zeta(\alpha, x_{\min})},

    where $\zeta(\alpha, x_{\min})$ is the Hurwitz zeta function.

    Parameters
    ----------
    alpha : float
      Exponent value of the distribution
    xmin : int
      Minimum value
    seed : integer, random_state, or None (default)
        Indicator of random number generation state.
        See :ref:`Randomness<randomness>`.

    Returns
    -------
    x : int
      Random value from Zipf distribution

    Raises
    ------
    ValueError:
      If xmin < 1 or
      If alpha <= 1

    Notes
    -----
    The rejection algorithm generates random values for a the power-law
    distribution in uniformly bounded expected time dependent on
    parameters.  See [1]_ for details on its operation.

    Examples
    --------
    >>> nx.zipf_rv(alpha=2, xmin=3, seed=42)  # doctest: +SKIP

    References
    ----------
    .. [1] Luc Devroye, Non-Uniform Random Variate Generation,
       Springer-Verlag, New York, 1986.
    """
    if xmin < 1:
        raise ValueError("xmin < 1")
    if alpha <= 1:
        raise ValueError("a <= 1.0")
    a1 = alpha - 1.0
    b = 2**a1
    while True:
        u = 1.0 - seed.random()  # u in (0,1]
        v = seed.random()  # v in [0,1)
        x = int(xmin * u**-(1.0 / a1))
        t = (1.0 + (1.0 / x))**a1
        if v * x * (t - 1.0) / (b - 1.0) <= t / b:
            break
    return x

def cumulative_distribution(distribution):
    """Returns normalized cumulative distribution from discrete distribution."""

    cdf = [0.0]
    psum = float(sum(distribution))
    for i in range(0, len(distribution)):
        cdf.append(cdf[i] + distribution[i] / psum)
    return cdf

@py_random_state(3)
def discrete_sequence(n, distribution=None, cdistribution=None, seed=None):
    """
    Return sample sequence of length n from a given discrete distribution
    or discrete cumulative distribution.

    One of the following must be specified.

    distribution = histogram of values, will be normalized

    cdistribution = normalized discrete cumulative distribution

    """
    import bisect

    if cdistribution is not None:
        cdf = cdistribution
    elif distribution is not None:
        cdf = cumulative_distribution(distribution)
    else:
        raise nx.NetworkXError(
            "discrete_sequence: distribution or cdistribution missing")

    # get a uniform random number
    inputseq = [seed.random() for i in range(n)]

    # choose from CDF
    seq = [bisect.bisect_left(cdf, s) - 1 for s in inputseq]
    return seq

@py_random_state(2)
def random_weighted_sample(mapping, k, seed=None):
    """Returns k items without replacement from a weighted sample.

    The input is a dictionary of items with weights as values.
    """
    if k > len(mapping):
        raise ValueError("sample larger than population")
    sample = set()

    return list(sample)

@py_random_state(1)
def weighted_choice(mapping, seed=None):
    """Returns a single element from a weighted sample.

    The input is a dictionary of items with weights as values.
    """
    # use roulette method
    rnd = seed.random() * sum(mapping.values())
    for k, w in mapping.items():
        rnd -= w
        if rnd < 0:
            return k


    ary = a
    a = astensor(a)

    if a.ndim < 2:
        raise ValueError("vsplit only works on tensors of 2 or more dimensions")
    return split(ary, indices_or_sections, 0)

from ..datasource import tensor as astensor
from .split import split

def vsplit(a, indices_or_sections):
    """
    Split a tensor into multiple sub-tensors vertically (row-wise).

    Please refer to the ``split`` documentation.  ``vsplit`` is equivalent
    to ``split`` with `axis=0` (default), the tensor is always split along the
    first axis regardless of the tensor dimension.

    See Also
    --------
    split : Split a tensor into multiple sub-tensors of equal size.

    Examples
    --------
    >>> import mars.tensor as mt

    >>> x = mt.arange(16.0).reshape(4, 4)
    >>> x.execute()
    array([[  0.,   1.,   2.,   3.],
           [  4.,   5.,   6.,   7.],
           [  8.,   9.,  10.,  11.],
           [ 12.,  13.,  14.,  15.]])
    >>> mt.vsplit(x, 2).execute()
    [array([[ 0.,  1.,  2.,  3.],
           [ 4.,  5.,  6.,  7.]]),
     array([[  8.,   9.,  10.,  11.],
           [ 12.,  13.,  14.,  15.]])]
    >>> mt.vsplit(x, mt.array([3, 6])).execute()
    [array([[  0.,   1.,   2.,   3.],
           [  4.,   5.,   6.,   7.],
           [  8.,   9.,  10.,  11.]]),
     array([[ 12.,  13.,  14.,  15.]]),
     array([], dtype=float64)]

    With a higher dimensional tensor the split is still along the first axis.

    >>> x = mt.arange(8.0).reshape(2, 2, 2)
    >>> x.execute()
    array([[[ 0.,  1.],
            [ 2.,  3.]],
           [[ 4.,  5.],
            [ 6.,  7.]]])
    >>> mt.vsplit(x, 2).execute()
    [array([[[ 0.,  1.],
            [ 2.,  3.]]]),
     array([[[ 4.,  5.],
            [ 6.,  7.]]])]

    """


    try:
      other = as_dimension(other)
    except (TypeError, ValueError):
      return NotImplemented
    if self._value is None or other.value is None:
      return None
    return self._value == other.value

"""Helper classes for tensor shape inference."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from tensorflow.core.framework import tensor_shape_pb2
from tensorflow.python.framework import dtypes
from tensorflow.python.util import compat
from tensorflow.python.util.tf_export import tf_export

@tf_export("Dimension")
class Dimension(object):
  """Represents the value of one dimension in a TensorShape."""

  def __init__(self, value):
    """Creates a new Dimension with the given value."""
    if value is None:
      self._value = None
    elif isinstance(value, dtypes.DType):
      raise TypeError("Cannot convert %s to Dimension" % value)
    else:
      self._value = int(value)
      if (not isinstance(value, compat.bytes_or_text_types) and
          self._value != value):
        raise ValueError("Ambiguous dimension: %s" % value)
      if self._value < 0:
        raise ValueError("Dimension %d must be >= 0" % self._value)

  def __repr__(self):
    return "Dimension(%s)" % repr(self._value)

  def __str__(self):
    value = self._value
    return "?" if value is None else str(value)

  def __eq__(self, other):
    """Returns true if `other` has the same known value as this Dimension."""

  def __ne__(self, other):
    """Returns true if `other` has a different known value from `self`."""
    try:
      other = as_dimension(other)
    except (TypeError, ValueError):
      return NotImplemented
    if self._value is None or other.value is None:
      return None
    return self._value != other.value

  def __int__(self):
    return self._value

  def __long__(self):
    return self._value

  def __index__(self):
    # Allow use in Python 3 range
    return self._value

  @property
  def value(self):
    """The value of this dimension, or None if it is unknown."""
    return self._value

  def is_compatible_with(self, other):
    """Returns true if `other` is compatible with this Dimension.

    Two known Dimensions are compatible if they have the same value.
    An unknown Dimension is compatible with all other Dimensions.

    Args:
      other: Another Dimension.

    Returns:
      True if this Dimension and `other` are compatible.
    """
    other = as_dimension(other)
    return (self._value is None or other.value is None or
            self._value == other.value)

  def assert_is_compatible_with(self, other):
    """Raises an exception if `other` is not compatible with this Dimension.

    Args:
      other: Another Dimension.

    Raises:
      ValueError: If `self` and `other` are not compatible (see
        is_compatible_with).
    """
    if not self.is_compatible_with(other):
      raise ValueError("Dimensions %s and %s are not compatible" % (self,
                                                                    other))

  def merge_with(self, other):
    """Returns a Dimension that combines the information in `self` and `other`.

    Dimensions are combined as follows:

    ```python
    tf.Dimension(n)   .merge_with(tf.Dimension(n))    == tf.Dimension(n)
    tf.Dimension(n)   .merge_with(tf.Dimension(None)) == tf.Dimension(n)
    tf.Dimension(None).merge_with(tf.Dimension(n))    == tf.Dimension(n)
    tf.Dimension(None).merge_with(tf.Dimension(None)) == tf.Dimension(None)
    tf.Dimension(n)   .merge_with(tf.Dimension(m))  # raises ValueError for n != m
    ```

    Args:
      other: Another Dimension.

    Returns:
      A Dimension containing the combined information of `self` and
      `other`.

    Raises:
      ValueError: If `self` and `other` are not compatible (see
        is_compatible_with).
    """
    other = as_dimension(other)
    self.assert_is_compatible_with(other)
    if self._value is None:
      return Dimension(other.value)
    else:
      return Dimension(self._value)

  def __add__(self, other):
    """Returns the sum of `self` and `other`.

    Dimensions are summed as follows:

    ```python
    tf.Dimension(m)    + tf.Dimension(n)    == tf.Dimension(m + n)
    tf.Dimension(m)    + tf.Dimension(None) == tf.Dimension(None)
    tf.Dimension(None) + tf.Dimension(n)    == tf.Dimension(None)
    tf.Dimension(None) + tf.Dimension(None) == tf.Dimension(None)
    ```

    Args:
      other: Another Dimension, or a value accepted by `as_dimension`.

    Returns:
      A Dimension whose value is the sum of `self` and `other`.
    """
    other = as_dimension(other)
    if self._value is None or other.value is None:
      return Dimension(None)
    else:
      return Dimension(self._value + other.value)

  def __radd__(self, other):
    """Returns the sum of `other` and `self`.

    Args:
      other: Another Dimension, or a value accepted by `as_dimension`.

    Returns:
      A Dimension whose value is the sum of `self` and `other`.
    """
    return self + other

  def __sub__(self, other):
    """Returns the subtraction of `other` from `self`.

    Dimensions are subtracted as follows:

    ```python
    tf.Dimension(m)    - tf.Dimension(n)    == tf.Dimension(m - n)
    tf.Dimension(m)    - tf.Dimension(None) == tf.Dimension(None)
    tf.Dimension(None) - tf.Dimension(n)    == tf.Dimension(None)
    tf.Dimension(None) - tf.Dimension(None) == tf.Dimension(None)
    ```

    Args:
      other: Another Dimension, or a value accepted by `as_dimension`.

    Returns:
      A Dimension whose value is the subtraction of `other` from `self`.
    """
    other = as_dimension(other)
    if self._value is None or other.value is None:
      return Dimension(None)
    else:
      return Dimension(self._value - other.value)

  def __rsub__(self, other):
    """Returns the subtraction of `self` from `other`.

    Args:
      other: Another Dimension, or a value accepted by `as_dimension`.

    Returns:
      A Dimension whose value is the subtraction of `self` from `other`.
    """
    other = as_dimension(other)
    if self._value is None or other.value is None:
      return Dimension(None)
    else:
      return Dimension(other.value - self._value)

  def __mul__(self, other):
    """Returns the product of `self` and `other`.

    Dimensions are summed as follows:

    ```python
    tf.Dimension(m)    * tf.Dimension(n)    == tf.Dimension(m * n)
    tf.Dimension(m)    * tf.Dimension(None) == tf.Dimension(None)
    tf.Dimension(None) * tf.Dimension(n)    == tf.Dimension(None)
    tf.Dimension(None) * tf.Dimension(None) == tf.Dimension(None)
    ```

    Args:
      other: Another Dimension, or a value accepted by `as_dimension`.

    Returns:
      A Dimension whose value is the product of `self` and `other`.
    """
    try:
      other = as_dimension(other)
    except (TypeError, ValueError):
      return NotImplemented

    if self._value is None or other.value is None:
      return Dimension(None)
    else:
      return Dimension(self._value * other.value)

  def __rmul__(self, other):
    """Returns the product of `self` and `other`.

    Args:
      other: Another Dimension, or a value accepted by `as_dimension`.

    Returns:
      A Dimension whose value is the product of `self` and `other`.
    """
    return self * other

  def __floordiv__(self, other):
    """Returns the quotient of `self` and `other` rounded down.

    Dimensions are divided as follows:

    ```python
    tf.Dimension(m)    // tf.Dimension(n)    == tf.Dimension(m // n)
    tf.Dimension(m)    // tf.Dimension(None) == tf.Dimension(None)
    tf.Dimension(None) // tf.Dimension(n)    == tf.Dimension(None)
    tf.Dimension(None) // tf.Dimension(None) == tf.Dimension(None)
    ```

    Args:
      other: Another Dimension, or a value accepted by `as_dimension`.

    Returns:
      A `Dimension` whose value is the integer quotient of `self` and `other`.
    """
    try:
      other = as_dimension(other)
    except (TypeError, ValueError):
      return NotImplemented
    if self._value is None or other.value is None:
      return Dimension(None)
    else:
      return Dimension(self._value // other.value)

  def __rfloordiv__(self, other):
    """Returns the quotient of `other` and `self` rounded down.

    Args:
      other: Another Dimension, or a value accepted by `as_dimension`.

    Returns:
      A `Dimension` whose value is the integer quotient of `self` and `other`.
    """
    other = as_dimension(other)
    if self._value is None or other.value is None:
      return Dimension(None)
    else:
      return Dimension(other.value // self._value)

  def __div__(self, other):
    """DEPRECATED: Use `__floordiv__` via `x // y` instead.

    This function exists only for backwards compatibility purposes; new code
    should use `__floordiv__` via the syntax `x // y`.  Using `x // y`
    communicates clearly that the result rounds down, and is forward compatible
    to Python 3.

    Args:
      other: Another `Dimension`.

    Returns:
      A `Dimension` whose value is the integer quotient of `self` and `other`.
    """
    return self // other

  def __mod__(self, other):
    """Returns `self` modulo `other`.

    Dimension moduli are computed as follows:

    ```python
    tf.Dimension(m)    % tf.Dimension(n)    == tf.Dimension(m % n)
    tf.Dimension(m)    % tf.Dimension(None) == tf.Dimension(None)
    tf.Dimension(None) % tf.Dimension(n)    == tf.Dimension(None)
    tf.Dimension(None) % tf.Dimension(None) == tf.Dimension(None)
    ```

    Args:
      other: Another Dimension, or a value accepted by `as_dimension`.

    Returns:
      A Dimension whose value is `self` modulo `other`.
    """
    try:
      other = as_dimension(other)
    except (TypeError, ValueError):
      return NotImplemented
    if self._value is None or other.value is None:
      return Dimension(None)
    else:
      return Dimension(self._value % other.value)

  def __rmod__(self, other):
    """Returns `other` modulo `self`.

    Args:
      other: Another Dimension, or a value accepted by `as_dimension`.

    Returns:
      A Dimension whose value is `other` modulo `self`.
    """
    try:
      other = as_dimension(other)
    except (TypeError, ValueError):
      return NotImplemented
    return other % self

  def __lt__(self, other):
    """Returns True if `self` is known to be less than `other`.

    Dimensions are compared as follows:

    ```python
    (tf.Dimension(m)    < tf.Dimension(n))    == (m < n)
    (tf.Dimension(m)    < tf.Dimension(None)) == None
    (tf.Dimension(None) < tf.Dimension(n))    == None
    (tf.Dimension(None) < tf.Dimension(None)) == None
    ```

    Args:
      other: Another Dimension.

    Returns:
      The value of `self.value < other.value` if both are known, otherwise
      None.
    """
    other = as_dimension(other)
    if self._value is None or other.value is None:
      return None
    else:
      return self._value < other.value

  def __le__(self, other):
    """Returns True if `self` is known to be less than or equal to `other`.

    Dimensions are compared as follows:

    ```python
    (tf.Dimension(m)    <= tf.Dimension(n))    == (m <= n)
    (tf.Dimension(m)    <= tf.Dimension(None)) == None
    (tf.Dimension(None) <= tf.Dimension(n))    == None
    (tf.Dimension(None) <= tf.Dimension(None)) == None
    ```

    Args:
      other: Another Dimension.

    Returns:
      The value of `self.value <= other.value` if both are known, otherwise
      None.
    """
    other = as_dimension(other)
    if self._value is None or other.value is None:
      return None
    else:
      return self._value <= other.value

  def __gt__(self, other):
    """Returns True if `self` is known to be greater than `other`.

    Dimensions are compared as follows:

    ```python
    (tf.Dimension(m)    > tf.Dimension(n))    == (m > n)
    (tf.Dimension(m)    > tf.Dimension(None)) == None
    (tf.Dimension(None) > tf.Dimension(n))    == None
    (tf.Dimension(None) > tf.Dimension(None)) == None
    ```

    Args:
      other: Another Dimension.

    Returns:
      The value of `self.value > other.value` if both are known, otherwise
      None.
    """
    other = as_dimension(other)
    if self._value is None or other.value is None:
      return None
    else:
      return self._value > other.value

  def __ge__(self, other):
    """Returns True if `self` is known to be greater than or equal to `other`.

    Dimensions are compared as follows:

    ```python
    (tf.Dimension(m)    >= tf.Dimension(n))    == (m >= n)
    (tf.Dimension(m)    >= tf.Dimension(None)) == None
    (tf.Dimension(None) >= tf.Dimension(n))    == None
    (tf.Dimension(None) >= tf.Dimension(None)) == None
    ```

    Args:
      other: Another Dimension.

    Returns:
      The value of `self.value >= other.value` if both are known, otherwise
      None.
    """
    other = as_dimension(other)
    if self._value is None or other.value is None:
      return None
    else:
      return self._value >= other.value

  def __reduce__(self):
    return Dimension, (self._value,)

def as_dimension(value):
  """Converts the given value to a Dimension.

  A Dimension input will be returned unmodified.
  An input of `None` will be converted to an unknown Dimension.
  An integer input will be converted to a Dimension with that value.

  Args:
    value: The value to be converted.

  Returns:
    A Dimension corresponding to the given value.
  """
  if isinstance(value, Dimension):
    return value
  else:
    return Dimension(value)

@tf_export("TensorShape")
class TensorShape(object):
  """Represents the shape of a `Tensor`.

  A `TensorShape` represents a possibly-partial shape specification for a
  `Tensor`. It may be one of the following:

  * *Fully-known shape:* has a known number of dimensions and a known size
    for each dimension. e.g. `TensorShape([16, 256])`
  * *Partially-known shape:* has a known number of dimensions, and an unknown
    size for one or more dimension. e.g. `TensorShape([None, 256])`
  * *Unknown shape:* has an unknown number of dimensions, and an unknown
    size in all dimensions. e.g. `TensorShape(None)`

  If a tensor is produced by an operation of type `"Foo"`, its shape
  may be inferred if there is a registered shape function for
  `"Foo"`. See [Shape
  functions](http://example.com)
  for details of shape functions and how to register them. Alternatively,
  the shape may be set explicitly using `tf.Tensor.set_shape`.
  """

  def __init__(self, dims):
    """Creates a new TensorShape with the given dimensions.

    Args:
      dims: A list of Dimensions, or None if the shape is unspecified.
        DEPRECATED: A single integer is treated as a singleton list.

    Raises:
      TypeError: If dims cannot be converted to a list of dimensions.
    """
    # TODO(irving): Eliminate the single integer special case.
    if dims is None:
      self._dims = None
    elif isinstance(dims, compat.bytes_or_text_types):
      raise TypeError("A string has ambiguous TensorShape, please wrap in a "
                      "list or convert to an int: %s" % dims)
    elif isinstance(dims, tensor_shape_pb2.TensorShapeProto):
      if dims.unknown_rank:
        self._dims = None
      else:
        self._dims = [
            # Protos store variable-size dimensions as -1
            as_dimension(dim.size if dim.size != -1 else None)
            for dim in dims.dim
        ]
    elif isinstance(dims, TensorShape):
      self._dims = dims.dims
    else:
      try:
        dims_iter = iter(dims)
      except TypeError:
        # Treat as a singleton dimension
        self._dims = [as_dimension(dims)]
      else:
        # Got a list of dimensions
        self._dims = [as_dimension(d) for d in dims_iter]
    self._ndims = None

  def __repr__(self):
    return "TensorShape(%r)" % self._dims

  def __str__(self):
    if self.ndims is None:
      return "<unknown>"
    elif self.ndims == 1:
      return "(%s,)" % self._dims[0]
    else:
      return "(%s)" % ", ".join(str(d) for d in self._dims)

  @property
  def dims(self):
    """Returns a list of Dimensions, or None if the shape is unspecified."""
    return self._dims

  @dims.setter
  def dims(self, dims):
    self._dims = dims
    self._ndims = None

  @property
  def ndims(self):
    """Returns the rank of this shape, or None if it is unspecified."""
    if self._dims is None:
      return None
    else:
      if self._ndims is None:
        self._ndims = len(self._dims)
      return self._ndims

  def __len__(self):
    """Returns the rank of this shape, or raises ValueError if unspecified."""
    if self._dims is None:
      raise ValueError("Cannot take the length of Shape with unknown rank.")
    return self.ndims

  def __bool__(self):
    """Returns True if this shape contains non-zero information."""
    return self._dims is not None

  # Python 3 wants __bool__, Python 2.7 wants __nonzero__
  __nonzero__ = __bool__

  def __iter__(self):
    """Returns `self.dims` if the rank is known, otherwise raises ValueError."""
    if self._dims is None:
      raise ValueError("Cannot iterate over a shape with unknown rank.")
    else:
      return iter(self._dims)

  def __getitem__(self, key):
    """Returns the value of a dimension or a shape, depending on the key.

    Args:
      key: If `key` is an integer, returns the dimension at that index;
        otherwise if `key` is a slice, returns a TensorShape whose
        dimensions are those selected by the slice from `self`.

    Returns:
      A dimension if `key` is an integer, or a `TensorShape` if `key` is a
      slice.

    Raises:
      ValueError: If `key` is a slice and `self` is completely unknown and
        the step is set.
    """
    if self._dims is not None:
      if isinstance(key, slice):
        return TensorShape(self._dims[key])
      else:
        return self._dims[key]
    else:
      if isinstance(key, slice):
        start = key.start if key.start is not None else 0
        stop = key.stop

        if key.step is not None:
          # TODO(mrry): Handle these maybe.
          raise ValueError("Steps are not yet handled")
        if stop is None:
          # NOTE(mrry): This implies that TensorShape(None) is compatible with
          # TensorShape(None)[1:], which is obviously not true. It would be
          # possible to track the number of dimensions symbolically,
          # and perhaps we should do that.
          return unknown_shape()
        elif start < 0 or stop < 0:
          # TODO(mrry): Handle this better, as it will be useful for handling
          # suffixes of otherwise unknown shapes.
          return unknown_shape()
        else:
          return unknown_shape(ndims=stop - start)
      else:
        return Dimension(None)

  def num_elements(self):
    """Returns the total number of elements, or none for incomplete shapes."""
    if self.is_fully_defined():
      size = 1
      for dim in self._dims:
        size *= dim.value
      return size
    else:
      return None

  def merge_with(self, other):
    """Returns a `TensorShape` combining the information in `self` and `other`.

    The dimensions in `self` and `other` are merged elementwise,
    according to the rules defined for `Dimension.merge_with()`.

    Args:
      other: Another `TensorShape`.

    Returns:
      A `TensorShape` containing the combined information of `self` and
      `other`.

    Raises:
      ValueError: If `self` and `other` are not compatible.
    """
    other = as_shape(other)
    if self._dims is None:
      return other
    else:
      try:
        self.assert_same_rank(other)
        new_dims = []
        for i, dim in enumerate(self._dims):
          new_dims.append(dim.merge_with(other[i]))
        return TensorShape(new_dims)
      except ValueError:
        raise ValueError("Shapes %s and %s are not compatible" % (self, other))

  def concatenate(self, other):
    """Returns the concatenation of the dimension in `self` and `other`.

    *N.B.* If either `self` or `other` is completely unknown,
    concatenation will discard information about the other shape. In
    future, we might support concatenation that preserves this
    information for use with slicing.

    Args:
      other: Another `TensorShape`.

    Returns:
      A `TensorShape` whose dimensions are the concatenation of the
      dimensions in `self` and `other`.
    """
    # TODO(mrry): Handle the case where we concatenate a known shape with a
    # completely unknown shape, so that we can use the partial information.
    other = as_shape(other)
    if self._dims is None or other.dims is None:
      return unknown_shape()
    else:
      return TensorShape(self._dims + other.dims)

  def assert_same_rank(self, other):
    """Raises an exception if `self` and `other` do not have compatible ranks.

    Args:
      other: Another `TensorShape`.

    Raises:
      ValueError: If `self` and `other` do not represent shapes with the
        same rank.
    """
    other = as_shape(other)
    if self.ndims is not None and other.ndims is not None:
      if self.ndims != other.ndims:
        raise ValueError("Shapes %s and %s must have the same rank" % (self,
                                                                       other))

  def assert_has_rank(self, rank):
    """Raises an exception if `self` is not compatible with the given `rank`.

    Args:
      rank: An integer.

    Raises:
      ValueError: If `self` does not represent a shape with the given `rank`.
    """
    if self.ndims not in (None, rank):
      raise ValueError("Shape %s must have rank %d" % (self, rank))

  def with_rank(self, rank):
    """Returns a shape based on `self` with the given rank.

    This method promotes a completely unknown shape to one with a
    known rank.

    Args:
      rank: An integer.

    Returns:
      A shape that is at least as specific as `self` with the given rank.

    Raises:
      ValueError: If `self` does not represent a shape with the given `rank`.
    """
    try:
      return self.merge_with(unknown_shape(ndims=rank))
    except ValueError:
      raise ValueError("Shape %s must have rank %d" % (self, rank))

  def with_rank_at_least(self, rank):
    """Returns a shape based on `self` with at least the given rank.

    Args:
      rank: An integer.

    Returns:
      A shape that is at least as specific as `self` with at least the given
      rank.

    Raises:
      ValueError: If `self` does not represent a shape with at least the given
        `rank`.
    """
    if self.ndims is not None and self.ndims < rank:
      raise ValueError("Shape %s must have rank at least %d" % (self, rank))
    else:
      return self

  def with_rank_at_most(self, rank):
    """Returns a shape based on `self` with at most the given rank.

    Args:
      rank: An integer.

    Returns:
      A shape that is at least as specific as `self` with at most the given
      rank.

    Raises:
      ValueError: If `self` does not represent a shape with at most the given
        `rank`.
    """
    if self.ndims is not None and self.ndims > rank:
      raise ValueError("Shape %s must have rank at most %d" % (self, rank))
    else:
      return self

  def is_compatible_with(self, other):
    """Returns True iff `self` is compatible with `other`.

    Two possibly-partially-defined shapes are compatible if there
    exists a fully-defined shape that both shapes can represent. Thus,
    compatibility allows the shape inference code to reason about
    partially-defined shapes. For example:

    * TensorShape(None) is compatible with all shapes.

    * TensorShape([None, None]) is compatible with all two-dimensional
      shapes, such as TensorShape([32, 784]), and also TensorShape(None). It is
      not compatible with, for example, TensorShape([None]) or
      TensorShape([None, None, None]).

    * TensorShape([32, None]) is compatible with all two-dimensional shapes
      with size 32 in the 0th dimension, and also TensorShape([None, None])
      and TensorShape(None). It is not compatible with, for example,
      TensorShape([32]), TensorShape([32, None, 1]) or TensorShape([64, None]).

    * TensorShape([32, 784]) is compatible with itself, and also
      TensorShape([32, None]), TensorShape([None, 784]), TensorShape([None,
      None]) and TensorShape(None). It is not compatible with, for example,
      TensorShape([32, 1, 784]) or TensorShape([None]).

    The compatibility relation is reflexive and symmetric, but not
    transitive. For example, TensorShape([32, 784]) is compatible with
    TensorShape(None), and TensorShape(None) is compatible with
    TensorShape([4, 4]), but TensorShape([32, 784]) is not compatible with
    TensorShape([4, 4]).

    Args:
      other: Another TensorShape.

    Returns:
      True iff `self` is compatible with `other`.

    """
    other = as_shape(other)
    if self._dims is not None and other.dims is not None:
      if self.ndims != other.ndims:
        return False
      for x_dim, y_dim in zip(self._dims, other.dims):
        if not x_dim.is_compatible_with(y_dim):
          return False
    return True

  def assert_is_compatible_with(self, other):
    """Raises exception if `self` and `other` do not represent the same shape.

    This method can be used to assert that there exists a shape that both
    `self` and `other` represent.

    Args:
      other: Another TensorShape.

    Raises:
      ValueError: If `self` and `other` do not represent the same shape.
    """
    if not self.is_compatible_with(other):
      raise ValueError("Shapes %s and %s are incompatible" % (self, other))

  def most_specific_compatible_shape(self, other):
    """Returns the most specific TensorShape compatible with `self` and `other`.

    * TensorShape([None, 1]) is the most specific TensorShape compatible with
      both TensorShape([2, 1]) and TensorShape([5, 1]). Note that
      TensorShape(None) is also compatible with above mentioned TensorShapes.

    * TensorShape([1, 2, 3]) is the most specific TensorShape compatible with
      both TensorShape([1, 2, 3]) and TensorShape([1, 2, 3]). There are more
      less specific TensorShapes compatible with above mentioned TensorShapes,
      e.g. TensorShape([1, 2, None]), TensorShape(None).

    Args:
      other: Another `TensorShape`.

    Returns:
      A `TensorShape` which is the most specific compatible shape of `self`
      and `other`.
    """

    other = as_shape(other)
    if self._dims is None or other.dims is None or self.ndims != other.ndims:
      return unknown_shape()

    dims = [(Dimension(None))] * self.ndims
    for i, (d1, d2) in enumerate(zip(self._dims, other.dims)):
      if d1 is not None and d2 is not None and d1 == d2:
        dims[i] = d1
    return TensorShape(dims)

  def is_fully_defined(self):
    """Returns True iff `self` is fully defined in every dimension."""
    return (self._dims is not None and all(dim.value is not None
                                           for dim in self._dims))

  def assert_is_fully_defined(self):
    """Raises an exception if `self` is not fully defined in every dimension.

    Raises:
      ValueError: If `self` does not have a known value for every dimension.
    """
    if not self.is_fully_defined():
      raise ValueError("Shape %s is not fully defined" % self)

  def as_list(self):
    """Returns a list of integers or `None` for each dimension.

    Returns:
      A list of integers or `None` for each dimension.

    Raises:
      ValueError: If `self` is an unknown shape with an unknown rank.
    """
    if self._dims is None:
      raise ValueError("as_list() is not defined on an unknown TensorShape.")
    return [dim.value for dim in self._dims]

  def as_proto(self):
    """Returns this shape as a `TensorShapeProto`."""
    if self._dims is None:
      return tensor_shape_pb2.TensorShapeProto(unknown_rank=True)
    else:
      return tensor_shape_pb2.TensorShapeProto(dim=[
          tensor_shape_pb2.TensorShapeProto.Dim(size=-1
                                                if d.value is None else d.value)
          for d in self._dims
      ])

  def __eq__(self, other):
    """Returns True if `self` is equivalent to `other`."""
    try:
      other = as_shape(other)
    except TypeError:
      return NotImplemented
    return self._dims == other.dims

  def __ne__(self, other):
    """Returns True if `self` is known to be different from `other`."""
    try:
      other = as_shape(other)
    except TypeError:
      return NotImplemented
    if self.ndims is None or other.ndims is None:
      raise ValueError("The inequality of unknown TensorShapes is undefined.")
    if self.ndims != other.ndims:
      return True
    return self._dims != other.dims

  def __reduce__(self):
    return TensorShape, (self._dims,)

def as_shape(shape):
  """Converts the given object to a TensorShape."""
  if isinstance(shape, TensorShape):
    return shape
  else:
    return TensorShape(shape)

def unknown_shape(ndims=None):
  """Returns an unknown TensorShape, optionally with a known rank.

  Args:
    ndims: (Optional) If specified, the number of dimensions in the shape.

  Returns:
    An unknown TensorShape.
  """
  if ndims is None:
    return TensorShape(None)
  else:
    return TensorShape([Dimension(None)] * ndims)

_SCALAR_SHAPE = TensorShape([])

def scalar():
  """Returns a shape representing a scalar."""
  return _SCALAR_SHAPE

def vector(length):
  """Returns a shape representing a vector.

  Args:
    length: The length of the vector, which may be None if unknown.

  Returns:
    A TensorShape representing a vector of the given length.
  """
  return TensorShape([length])

def matrix(rows, cols):
  """Returns a shape representing a matrix.

  Args:
    rows: The number of rows in the matrix, which may be None if unknown.
    cols: The number of columns in the matrix, which may be None if unknown.

  Returns:
    A TensorShape representing a matrix of the given size.
  """
  return TensorShape([rows, cols])